DecodeConfig

Author: bep

common/hugio/writers.go

@@ -15,8 +15,6 @@ package hugio
 
 import (
 	"io"
-
-	"github.com/gohugoio/hugo/common/hdebug"
 )
 
 // As implemented by strings.Builder.
@@ -127,18 +125,3 @@ func (c PipeReadWriteCloser) Close() (err error) {
 func (c PipeReadWriteCloser) WriteString(s string) (int, error) {
 	return c.PipeWriter.Write([]byte(s))
 }
-
-// CounterWriter is an io.Writer that counts the number of bytes written to it.
-// and prints a message every 1024 bytes.
-type CounterWriter struct {
-	count int64
-}
-
-func (cw *CounterWriter) Write(p []byte) (n int, err error) {
-	n = len(p)
-	cw.count += int64(n)
-	if cw.count%1024 == 0 {
-		hdebug.Printf("%d bytes written", cw.count)
-	}
-	return n, nil
-}

common/maps/map.go

@@ -74,10 +74,10 @@ func (m *Map[K, T]) Set(key K, value T) {
 }
 
 // WithWriteLock executes the given function with a write lock on the map.
-func (m *Map[K, T]) WithWriteLock(f func(m map[K]T)) {
+func (m *Map[K, T]) WithWriteLock(f func(m map[K]T) error) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
-	f(m.m)
+	return f(m.m)
 }
 
 // SetIfAbsent sets the given key to the given value if the key does not already exist in the map.

common/maps/map_test.go

@@ -62,8 +62,9 @@ func TestMap(t *testing.T) {
 	c.Assert(found, qt.Equals, true)
 	c.Assert(v, qt.Equals, 300)
 
-	m.WithWriteLock(func(m map[string]int) {
+	m.WithWriteLock(func(m map[string]int) error {
 		m["f"] = 500
+		return nil
 	})
 	v, found = m.Lookup("f")
 	c.Assert(found, qt.Equals, true)

hugolib/content_map_page_assembler.go

@@ -282,8 +282,9 @@ func (a *allPagesAssembler) doCreatePages(prefix string, depth int) error {
 						default:
 							// Skip this page.
 							a.droppedPages.WithWriteLock(
-								func(m map[*Site][]string) {
+								func(m map[*Site][]string) error {
 									m[site] = append(m[site], s)
+									return nil
 								},
 							)
 
@@ -635,15 +636,16 @@ func (a *allPagesAssembler) doCreatePages(prefix string, depth int) error {
 								continue
 							}
 							t := term{view: viewName, term: v}
-							a.seenTerms.WithWriteLock(func(m map[term]sitesmatrix.Vectors) {
+							a.seenTerms.WithWriteLock(func(m map[term]sitesmatrix.Vectors) error {
 								vectors, found := m[t]
 								if !found {
 									m[t] = sitesmatrix.Vectors{
 										vec: struct{}{},
 									}
-									return
+									return nil
 								}
 								vectors[vec] = struct{}{}
+								return nil
 							})
 						}
 					}

internal/warpc/warpc.go

@@ -33,7 +33,9 @@ import (
 	"github.com/bep/textandbinaryreader"
 
 	"github.com/gohugoio/hugo/common/hdebug"
+	"github.com/gohugoio/hugo/common/hstrings"
 	"github.com/gohugoio/hugo/common/hugio"
+	"github.com/gohugoio/hugo/common/maps"
 	"golang.org/x/sync/errgroup"
 
 	"github.com/tetratelabs/wazero"
@@ -42,6 +44,11 @@ import (
 	"github.com/tetratelabs/wazero/imports/wasi_snapshot_preview1"
 )
 
+const (
+	MessageKindJSON string = "json"
+	MessageKindBlob string = "blob"
+)
+
 const currentVersion = 1
 
 //go:embed wasm/quickjs.wasm
@@ -59,13 +66,51 @@ type Header struct {
 	// Command is the command to execute.
 	Command string `json:"command"`
 
+	// RequestKinds is a list of kinds in this RPC request,
+	// e.g. {"json", "blob"}, or {"json"}.
+	RequestKinds []string `json:"requestKinds"`
+	// ResponseKinds is a list of kinds expected in the response,
+	// e.g. {"json", "blob"}, or {"json"}.
+	ResponseKinds []string `json:"responseKinds"`
+
 	// Set in the response if there was an error.
 	Err string `json:"err"`
 
 	// Warnings is a list of warnings that may be returned in the response.
 	Warnings []string `json:"warnings,omitempty"`
 }
 
+func (m *Header) init() error {
+	if m.ID == 0 {
+		return errors.New("ID must not be 0 (note that this must be unique within the current request set time window)")
+	}
+	if m.Version == 0 {
+		m.Version = currentVersion
+	}
+	if len(m.RequestKinds) == 0 {
+		m.RequestKinds = []string{string(MessageKindJSON)}
+	}
+	if len(m.ResponseKinds) == 0 {
+		m.ResponseKinds = []string{string(MessageKindJSON)}
+	}
+	if m.Version != currentVersion {
+		return fmt.Errorf("unsupported version: %d", m.Version)
+	}
+	for range 2 {
+		if len(m.RequestKinds) > 2 {
+			return fmt.Errorf("invalid number of request kinds: %d", len(m.RequestKinds))
+		}
+		if len(m.ResponseKinds) > 2 {
+			return fmt.Errorf("invalid number of response kinds: %d", len(m.ResponseKinds))
+		}
+		m.RequestKinds = hstrings.UniqueStringsReuse(m.RequestKinds)
+		m.ResponseKinds = hstrings.UniqueStringsReuse(m.ResponseKinds)
+
+	}
+
+	return nil
+}
+
 type Message[T any] struct {
 	Header Header `json:"header"`
 	Data   T      `json:"data"`
@@ -75,6 +120,10 @@ func (m Message[T]) GetID() uint32 {
 	return m.Header.ID
 }
 
+func (m *Message[T]) init() error {
+	return m.Header.init()
+}
+
 type SourceProvider interface {
 	GetSource() io.Reader
 	GetSourceLength() uint32
@@ -160,9 +209,6 @@ func putTimer(t *time.Timer) {
 // Execute sends a request to the dispatcher and waits for the response.
 func (p *dispatcherPool[Q, R]) Execute(ctx context.Context, q Message[Q]) (Message[R], error) {
 	d := p.getDispatcher()
-	if q.GetID() == 0 {
-		return d.zeroR, errors.New("ID must not be 0 (note that this must be unique within the current request set time window)")
-	}
 
 	call, err := d.newCall(q)
 	if err != nil {
@@ -195,20 +241,22 @@ func (p *dispatcherPool[Q, R]) Execute(ctx context.Context, q Message[Q]) (Messa
 	if err == nil && resp.Header.Err != "" {
 		err = errors.New(resp.Header.Err)
 	}
+
 	return resp, err
 }
 
 func (d *dispatcher[Q, R]) newCall(q Message[Q]) (*call[Q, R], error) {
-	responseCountdown := &atomic.Int32{}
-	responseCountdown.Add(1) // Default is JSON response only.
-	if _, ok := any(d.zeroQ.Data).(DestinationProvider); ok {
-		// Expecting JSON followed by binary blob.
-		responseCountdown.Add(1)
+	if err := q.init(); err != nil {
+		return nil, err
+	}
+	responseKinds := maps.NewMap[string, bool]()
+	for _, rk := range q.Header.ResponseKinds {
+		responseKinds.Set(rk, true)
 	}
 	call := &call[Q, R]{
-		donec:             make(chan *call[Q, R], 1),
-		request:           q,
-		responseCountdown: responseCountdown,
+		donec:         make(chan *call[Q, R], 1),
+		request:       q,
+		responseKinds: responseKinds,
 	}
 
 	if d.shutdown || d.closing {
@@ -240,7 +288,7 @@ func (d *dispatcher[Q, R]) send(call *call[Q, R]) error {
 	}
 	if sp, ok := any(call.request.Data).(SourceProvider); ok {
 		if sp.GetSourceLength() == 0 {
-			panic("source length must be greater than 0")
+			panic(fmt.Sprintf("source length must be greater than 0, header: %+v", call.request.Header))
 		}
 
 		hdebug.Printf(" == == === Blob header %d %d", call.request.GetID(), sp.GetSourceLength())
@@ -272,7 +320,7 @@ func (d *dispatcher[Q, R]) inputBlobs() {
 
 		call := d.pendingCall(id)
 
-		hdebug.Printf("START === === === Read blob header id: %d len: %d countdown: %d", id, length, call.responseCountdown.Load())
+		hdebug.Printf("START === === === Read blob header id: %d len: %d ", id, length)
 
 		if err := call.handleBlob(lr); err != nil {
 			inputErr = err
@@ -282,18 +330,27 @@ func (d *dispatcher[Q, R]) inputBlobs() {
 			inputErr = fmt.Errorf("blob %d: expected to read %d more bytes", id, lr.N)
 			break
 		}
-		call.responseCountdown.Add(-1)
-		if call.responseCountdown.Load() <= 0 {
-			d.mu.Lock()
-			delete(d.pending, id)
-			d.mu.Unlock()
-			call.done()
+		if err := call.responseKinds.WithWriteLock(
+			func(m map[string]bool) error {
+				if _, ok := m[MessageKindBlob]; !ok {
+					return fmt.Errorf("unexpected blob response for %q call ID %d", call.request.Header.Command, id)
+				}
+				delete(m, MessageKindBlob)
+				// Message exchange is complete.
+				d.mu.Lock()
+				delete(d.pending, id)
+				d.mu.Unlock()
+				call.done()
+				return nil
+			}); err != nil {
+			inputErr = err
+			break
 		}
-		hdebug.Printf("END === === === Read blob header id: %d len: %d countdown: %d", id, length, call.responseCountdown.Load())
+		hdebug.Printf("END === === === Read blob header id: %d len: %d", id, length)
 	}
 
-	if inputErr != nil && inputErr != io.EOF && inputErr != io.ErrClosedPipe {
-		fmt.Printf("ERR %s", inputErr) // TODO1
+	if inputErr != nil {
+		// panic(inputErr) // TODO1 fix me, consolidate with JSON error handling.
 	}
 }
 
@@ -309,17 +366,27 @@ func (d *dispatcher[Q, R]) inputJSON() {
 
 		call := d.pendingCall(r.GetID())
 
-		hdebug.Printf("END === === === get JSON id: %d  countdown: %d", r.GetID(), call.responseCountdown.Load())
-
-		call.responseCountdown.Add(-1)
-		call.response = r
-		if call.responseCountdown.Load() <= 0 || r.Header.Err != "" {
-			d.mu.Lock()
-			delete(d.pending, r.GetID())
-			d.mu.Unlock()
-			call.done() // TODO1 check that this can be called multiple times safely.
+		if err := call.responseKinds.WithWriteLock(
+			func(m map[string]bool) error {
+				call.response = r
+				if _, ok := m[MessageKindJSON]; !ok {
+					return fmt.Errorf("unexpected JSON response for call ID %d", r.GetID())
+				}
+				delete(m, MessageKindJSON)
+				if len(m) == 0 {
+					// Message exchange is complete.
+					d.mu.Lock()
+					delete(d.pending, r.GetID())
+					d.mu.Unlock()
+					call.done()
+				}
+				return nil
+			}); err != nil {
+			inputErr = err
+			break
 		}
-		hdebug.Printf("END === === === get JSON id: %d  countdown: %d", r.GetID(), call.responseCountdown.Load())
+
+		hdebug.Printf("END === === === get JSON id: %d", r.GetID())
 
 	}
 
@@ -355,11 +422,11 @@ func (d *dispatcher[Q, R]) pendingCall(id uint32) *call[Q, R] {
 }
 
 type call[Q, R any] struct {
-	request           Message[Q]
-	response          Message[R]
-	responseCountdown *atomic.Int32
-	err               error
-	donec             chan *call[Q, R]
+	request       Message[Q]
+	response      Message[R]
+	responseKinds *maps.Map[string, bool]
+	err           error
+	donec         chan *call[Q, R]
 }
 
 func (c *call[Q, R]) handleBlob(r io.Reader) error {