Handle Execute limit correctly by sending PortalSuspended

Author: JelteF

cache.go

@@ -2,12 +2,21 @@ package wire
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"iter"
 	"sync"
 
 	"github.com/jeroenrinzema/psql-wire/pkg/buffer"
+	"github.com/jeroenrinzema/psql-wire/pkg/types"
 )
 
+// Limit represents the maximum number of rows to be written.
+// Zero denotes "no limit".
+type Limit uint32
+
+const NoLimit Limit = 0
+
 type Statement struct {
 	fn         PreparedStatementFn
 	parameters []uint32
@@ -73,6 +82,112 @@ type Portal struct {
 	statement  *Statement
 	parameters []Parameter
 	formats    []FormatCode
+
+	// The iterator state (created by iter.Pull)
+	next func() (struct{}, bool)
+	stop func()
+	// Filled in by dataWriter.Complete when the handler finishes. Used to
+	// return the tag when re-executing a completed portal.
+	tag string
+	err error
+	// Set to true when execution of the portal has finished.
+	done bool
+
+	// pending is closed when the most recently launched async goroutine
+	// finishes. A new goroutine for the same portal waits on this channel
+	// before starting, so same-portal executes serialize while different
+	// portals run in parallel.
+	pending chan struct{}
+}
+
+// Close tears down the portal from outside the execution goroutine (e.g.
+// explicit Close message or portal re-bind). If an async goroutine is still
+// running, cleanup is deferred to a background goroutine that waits for it.
+func (p *Portal) Close() {
+	if p.pending != nil {
+		pending := p.pending
+		go func() {
+			<-pending
+			p.close()
+		}()
+		return
+	}
+	p.close()
+}
+
+// close tears down the iterator inline. Used from within execute where we
+// are already inside the execution context and don't need to wait on pending.
+func (p *Portal) close() {
+	if p.stop != nil {
+		p.stop()
+		p.next = nil
+		p.stop = nil
+	}
+}
+
+func portalSuspended(writer *buffer.Writer) error {
+	writer.Start(types.ServerPortalSuspended)
+	return writer.End()
+}
+
+func (p *Portal) execute(ctx context.Context, limit Limit, reader *buffer.Reader, writer *buffer.Writer) error {
+	if p.done {
+		// Re-executing an already completed portal simply returns the tag or
+		// error again.
+		if p.err != nil {
+			return p.err
+		}
+		return commandComplete(writer, p.tag)
+	}
+
+	if p.next == nil {
+		// This is the first execute call on this portal. So let's start the
+		// execution. Otherwise we continue from where we left off.
+		session, _ := GetSession(ctx)
+		// Create a simple push-style iterator (iter.Seq) around the
+		// statement.fn.
+		seq := func(yield func(struct{}) bool) {
+			dw := &dataWriter{
+				ctx:     ctx,
+				session: session,
+				columns: p.statement.columns,
+				formats: p.formats,
+				reader:  reader,
+				client:  writer,
+				yield:   yield,
+				tag:     &p.tag,
+			}
+			err := p.statement.fn(ctx, dw, p.parameters)
+			if err != nil && !errors.Is(err, ErrSuspendedHandlerClosed) {
+				p.err = err
+			}
+		}
+
+		// Then we convert that push-style iterator into a pull-style iterator,
+		// so we can suspend the iterator when we reach the row limit.
+		p.next, p.stop = iter.Pull(seq)
+	}
+
+	var count Limit
+	for {
+		if limit != NoLimit && count >= limit {
+			// We've reached the row limit. Suspend the portal and let the
+			// client know, so it can either issue a new Execute to continue or
+			// close the portal.
+			return portalSuspended(writer)
+		}
+
+		_, ok := p.next()
+		if !ok {
+			// The handler has finished. CommandComplete was already written
+			// by dataWriter.Complete.
+			p.close()
+			p.done = true
+			return p.err
+		}
+
+		count++
+	}
 }
 
 func DefaultPortalCacheFn() PortalCache {
@@ -92,6 +207,10 @@ func (cache *DefaultPortalCache) Bind(ctx context.Context, name string, stmt *St
 		cache.portals = map[string]*Portal{}
 	}
 
+	if existing, ok := cache.portals[name]; ok {
+		existing.Close()
+	}
+
 	cache.portals[name] = &Portal{
 		statement:  stmt,
 		parameters: parameters,
@@ -137,13 +256,15 @@ func (cache *DefaultPortalCache) Execute(ctx context.Context, name string, limit
 		return nil
 	}
 
-	session, _ := GetSession(ctx)
-	return portal.statement.fn(ctx, NewDataWriter(ctx, session, portal.statement.columns, portal.formats, limit, reader, writer), portal.parameters)
+	return portal.execute(ctx, limit, reader, writer)
 }
 
 func (cache *DefaultPortalCache) Delete(ctx context.Context, name string) error {
 	cache.mu.Lock()
 	defer cache.mu.Unlock()
+	if portal, ok := cache.portals[name]; ok {
+		portal.Close()
+	}
 	delete(cache.portals, name)
 	return nil
 }
@@ -153,10 +274,17 @@ func (cache *DefaultPortalCache) DeleteByStatement(ctx context.Context, stmt *St
 	defer cache.mu.Unlock()
 	for name, portal := range cache.portals {
 		if portal.statement == stmt {
+			portal.Close()
 			delete(cache.portals, name)
 		}
 	}
 	return nil
 }
 
-func (cache *DefaultPortalCache) Close() {}
+func (cache *DefaultPortalCache) Close() {
+	cache.mu.Lock()
+	defer cache.mu.Unlock()
+	for _, portal := range cache.portals {
+		portal.Close()
+	}
+}

command.go

@@ -1,6 +1,7 @@
 package wire
 
 import (
+	"bytes"
 	"context"
 	"errors"
 	"fmt"
@@ -56,6 +57,7 @@ type Session struct {
 	Statements StatementCache
 	Portals    PortalCache
 	Attributes map[string]interface{}
+	reader     *buffer.Reader
 
 	// pipelining
 	ParallelPipeline ParallelPipelineConfig
@@ -93,6 +95,7 @@ func isExtendedQueryMessage(t types.ClientMessage) bool {
 // This method keeps consuming messages until the client issues a close message
 // or the connection is terminated.
 func (srv *Session) consumeCommands(ctx context.Context, conn net.Conn, reader *buffer.Reader, writer *buffer.Writer) error {
+	srv.reader = reader
 	srv.logger.Debug("ready for query... starting to consume commands")
 
 	err := readyForQuery(writer, types.ServerIdle)
@@ -326,7 +329,13 @@ func (srv *Session) handleSimpleQuery(ctx context.Context, reader *buffer.Reader
 			return srv.WriteError(writer, err)
 		}
 
-		err = statements[index].fn(ctx, NewDataWriter(ctx, srv, statements[index].columns, nil, NoLimit, reader, writer), nil)
+		portal := &Portal{
+			statement: &Statement{
+				fn:      statements[index].fn,
+				columns: statements[index].columns,
+			},
+		}
+		err = portal.execute(ctx, NoLimit, reader, writer)
 		if err != nil {
 			return srv.WriteError(writer, err)
 		}
@@ -753,7 +762,11 @@ func (srv *Session) handleClose(ctx context.Context, reader *buffer.Reader, writ
 	return writer.End()
 }
 
-// executePipelined handles Execute in parallel pipeline mode
+// executePipelined handles Execute in parallel pipeline mode. When the
+// Execute carries a row limit (or the portal already has a suspended
+// iterator) it is queued as a synchronous event that runs inline during
+// drain, because it needs the portal's iter.Pull2 state. Otherwise it
+// launches a goroutine for parallel execution.
 func (srv *Session) executePipelined(ctx context.Context, writer *buffer.Writer, name string, limit uint32) error {
 	portal, err := srv.Portals.Get(ctx, name)
 	if err != nil {
@@ -764,42 +777,46 @@ func (srv *Session) executePipelined(ctx context.Context, writer *buffer.Writer,
 		return srv.drainQueueAndWriteError(ctx, writer, errors.New("unknown portal"))
 	}
 
-	// Create result channel and queue the event
-	resultChan := make(chan *QueuedDataWriter, 1)
-	srv.ResponseQueue.Enqueue(NewExecuteEvent(resultChan, portal.formats))
+	resultChan := make(chan *executeResult, 1)
+	srv.ResponseQueue.Enqueue(NewExecuteEvent(resultChan))
 
-	// Launch async execution
-	go srv.executeAsync(ctx, portal, limit, resultChan)
+	prev := portal.pending
+	done := make(chan struct{})
+	portal.pending = done
+	go srv.executeAsync(ctx, done, portal, Limit(limit), prev, resultChan)
 
 	return nil
 }
 
-// executeAsync runs the portal execution in a separate goroutine
-func (srv *Session) executeAsync(ctx context.Context, portal *Portal, limit uint32, resultChan chan<- *QueuedDataWriter) {
+// executeAsync runs portal.execute in a separate goroutine, capturing the
+// wire output into a buffer. If prev is non-nil, it waits for the previous
+// goroutine on the same portal to finish before starting.
+func (srv *Session) executeAsync(ctx context.Context, done chan struct{}, portal *Portal, limit Limit, prev chan struct{}, resultChan chan<- *executeResult) {
+	defer close(done)
 	defer close(resultChan)
 	defer func() {
 		if r := recover(); r != nil {
-			collector := NewQueuedDataWriter(ctx, portal.statement.columns, Limit(limit))
-			collector.SetError(fmt.Errorf("panic during execution: %v", r))
-			resultChan <- collector
+			resultChan <- &executeResult{err: fmt.Errorf("panic during execution: %v", r)}
 		}
 	}()
 
-	srv.logger.Debug("starting async execution",
-		slog.Bool("has_function", portal.statement.fn != nil))
+	if prev != nil {
+		<-prev
+	}
 
-	collector := NewQueuedDataWriter(ctx, portal.statement.columns, Limit(limit))
+	srv.logger.Debug("starting async execution")
 
-	err := portal.statement.fn(ctx, collector, portal.parameters)
-	if err != nil {
-		collector.SetError(err)
-	}
+	buf := &bytes.Buffer{}
+	w := buffer.NewWriter(srv.logger, buf)
+
+	err := portal.execute(ctx, limit, srv.reader, w)
+
+	result := &executeResult{buf: buf, err: err}
 
 	srv.logger.Debug("async execution complete",
-		slog.Bool("has_error", collector.GetError() != nil),
-		slog.Int("rows", int(collector.Written())))
+		slog.Bool("has_error", err != nil))
 
-	resultChan <- collector
+	resultChan <- result
 }
 
 // handleSync handles the Sync message (extended query protocol)
@@ -924,22 +941,16 @@ func (srv *Session) writeQueuedResponse(ctx context.Context, writer *buffer.Writ
 		return srv.writeColumnDescription(ctx, writer, event.Formats, event.Columns)
 
 	case ResponseExecute:
-		// Execute writes DataRows followed by CommandComplete
 		if event.Result == nil {
-			// No result yet, this shouldn't happen in normal flow
 			return errors.New("execute event has no result")
 		}
 
-		// Check for execution error
-		if err := event.Result.GetError(); err != nil {
-			return srv.WriteError(writer, err)
+		if event.Result.err != nil {
+			return srv.WriteError(writer, event.Result.err)
 		}
 
-		// Use DataWriter for correct encoding
-		// Note: We use NoLimit here because the result is already limited during execution
-		dataWriter := NewDataWriter(ctx, srv, event.Result.Columns(), event.Formats, NoLimit, nil, writer)
-
-		return event.Result.Replay(ctx, dataWriter)
+		_, err := writer.Write(event.Result.buf.Bytes())
+		return err
 
 	default:
 		return fmt.Errorf("unknown response event kind: %v", event.Kind)