Create a generic latch.go and make schemas use that (#3353)

Co-authored-by: Philip Dubé <[email protected]>

Author: iskakaushik

flow/cmd/handler.go

@@ -22,6 +22,7 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/internal"
 	"github.com/PeerDB-io/peerdb/flow/model"
 	"github.com/PeerDB-io/peerdb/flow/shared"
+	"github.com/PeerDB-io/peerdb/flow/shared/concurrency"
 	"github.com/PeerDB-io/peerdb/flow/shared/exceptions"
 	peerflow "github.com/PeerDB-io/peerdb/flow/workflows"
 )
@@ -264,14 +265,14 @@ func (h *FlowRequestHandler) shutdownFlow(
 	cancelCtx, cancel := context.WithTimeout(ctx, 5*time.Minute)
 	defer cancel()
 
-	errChan := make(chan error, 1)
+	errLatch := concurrency.NewLatch[error]()
 	go func() {
-		errChan <- dropFlowHandle.Get(cancelCtx, nil)
+		errLatch.Set(dropFlowHandle.Get(cancelCtx, nil))
 	}()
 
 	select {
-	case err := <-errChan:
-		if err != nil {
+	case <-errLatch.Chan():
+		if err := errLatch.Wait(); err != nil {
 			slog.Error("DropFlow workflow did not execute successfully", logs, slog.Any("error", err))
 			return fmt.Errorf("DropFlow workflow did not execute successfully: %w", err)
 		}
@@ -385,19 +386,17 @@ func (h *FlowRequestHandler) FlowStateChange(
 }
 
 func (h *FlowRequestHandler) handleCancelWorkflow(ctx context.Context, workflowID, runID string) error {
-	errChan := make(chan error, 1)
-
 	ctxWithTimeout, cancel := context.WithTimeout(ctx, 2*time.Minute)
 	defer cancel()
 
+	errLatch := concurrency.NewLatch[error]()
 	go func() {
-		err := h.temporalClient.CancelWorkflow(ctxWithTimeout, workflowID, runID)
-		errChan <- err
+		errLatch.Set(h.temporalClient.CancelWorkflow(ctxWithTimeout, workflowID, runID))
 	}()
 
 	select {
-	case err := <-errChan:
-		if err != nil {
+	case <-errLatch.Chan():
+		if err := errLatch.Wait(); err != nil {
 			slog.Error(fmt.Sprintf("unable to cancel PeerFlow workflow: %s. Attempting to terminate.", err.Error()))
 			terminationReason := fmt.Sprintf("workflow %s did not cancel in time.", workflowID)
 			if err := h.temporalClient.TerminateWorkflow(ctx, workflowID, runID, terminationReason); err != nil {

flow/connectors/postgres/sink_pg.go

@@ -12,13 +12,12 @@ import (
 	"github.com/PeerDB-io/peerdb/flow/connectors/postgres/sanitize"
 	"github.com/PeerDB-io/peerdb/flow/connectors/utils"
 	"github.com/PeerDB-io/peerdb/flow/shared"
+	"github.com/PeerDB-io/peerdb/flow/shared/concurrency"
 )
 
 type PgCopyShared struct {
-	schemaLatch chan struct{}
+	schemaLatch *concurrency.Latch[[]string]
 	err         error
-	schema      []string
-	schemaSet   bool
 }
 
 type PgCopyWriter struct {
@@ -33,17 +32,13 @@ type PgCopyReader struct {
 
 func NewPgCopyPipe() (PgCopyReader, PgCopyWriter) {
 	read, write := io.Pipe()
-	schema := PgCopyShared{schemaLatch: make(chan struct{})}
+	schema := PgCopyShared{schemaLatch: concurrency.NewLatch[[]string]()}
 	return PgCopyReader{PipeReader: read, schema: &schema},
 		PgCopyWriter{PipeWriter: write, schema: &schema}
 }
 
 func (p PgCopyWriter) SetSchema(schema []string) {
-	if !p.schema.schemaSet {
-		p.schema.schema = schema
-		close(p.schema.schemaLatch)
-		p.schema.schemaSet = true
-	}
+	p.schema.schemaLatch.Set(schema)
 }
 
 func (p PgCopyWriter) ExecuteQueryWithTx(
@@ -109,8 +104,7 @@ func (p PgCopyWriter) Close(err error) {
 }
 
 func (p PgCopyReader) GetColumnNames() ([]string, error) {
-	<-p.schema.schemaLatch
-	return p.schema.schema, p.schema.err
+	return p.schema.schemaLatch.Wait(), p.schema.err
 }
 
 func (p PgCopyReader) CopyInto(ctx context.Context, c *PostgresConnector, tx pgx.Tx, table pgx.Identifier) (int64, error) {

flow/model/qrecord_stream.go

@@ -1,46 +1,38 @@
 package model
 
 import (
+	"github.com/PeerDB-io/peerdb/flow/shared/concurrency"
 	"github.com/PeerDB-io/peerdb/flow/shared/types"
 )
 
 type QRecordStream struct {
-	schemaLatch chan struct{}
+	schemaLatch *concurrency.Latch[types.QRecordSchema]
 	Records     chan []types.QValue
 	err         error
-	schema      types.QRecordSchema
-	schemaSet   bool
 }
 
 func NewQRecordStream(buffer int) *QRecordStream {
 	return &QRecordStream{
-		schemaLatch: make(chan struct{}),
+		schemaLatch: concurrency.NewLatch[types.QRecordSchema](),
 		Records:     make(chan []types.QValue, buffer),
-		schema:      types.QRecordSchema{},
 		err:         nil,
-		schemaSet:   false,
 	}
 }
 
 func (s *QRecordStream) Schema() (types.QRecordSchema, error) {
-	<-s.schemaLatch
-	return s.schema, s.Err()
+	return s.schemaLatch.Wait(), s.Err()
 }
 
 func (s *QRecordStream) SetSchema(schema types.QRecordSchema) {
-	if !s.schemaSet {
-		s.schema = schema
-		s.schemaSet = true
-		close(s.schemaLatch)
-	}
+	s.schemaLatch.Set(schema)
 }
 
 func (s *QRecordStream) IsSchemaSet() bool {
-	return s.schemaSet
+	return s.schemaLatch.IsSet()
 }
 
 func (s *QRecordStream) SchemaChan() <-chan struct{} {
-	return s.schemaLatch
+	return s.schemaLatch.Chan()
 }
 
 func (s *QRecordStream) Err() error {
@@ -55,7 +47,7 @@ func (s *QRecordStream) Close(err error) {
 		s.err = err
 		close(s.Records)
 	}
-	if !s.schemaSet {
+	if !s.schemaLatch.IsSet() {
 		s.SetSchema(types.QRecordSchema{})
 	}
 }

flow/shared/concurrency/latch.go

@@ -0,0 +1,47 @@
+package concurrency
+
+import "sync"
+
+// Latch is a thread-safe value holder that can be set once and read many times.
+// Once Set is called, all future Wait calls will immediately return the set value.
+type Latch[T any] struct {
+	val   T
+	ready chan struct{}
+	once  sync.Once
+}
+
+// NewLatch creates a new Latch for type T.
+func NewLatch[T any]() *Latch[T] {
+	return &Latch[T]{ready: make(chan struct{})}
+}
+
+// Set publishes the value exactly once; later calls are no-ops.
+func (l *Latch[T]) Set(v T) {
+	l.once.Do(func() {
+		l.val = v
+		close(l.ready) // broadcasts to all waiters
+	})
+}
+
+// Wait blocks until Set is called, then returns the value.
+// Safe to call many times; subsequent calls return immediately.
+func (l *Latch[T]) Wait() T {
+	<-l.ready
+	return l.val
+}
+
+// Chan returns a channel that will be closed when the value is set.
+// Useful for select statements.
+func (l *Latch[T]) Chan() <-chan struct{} {
+	return l.ready
+}
+
+// IsSet returns true if the value has been set.
+func (l *Latch[T]) IsSet() bool {
+	select {
+	case <-l.ready:
+		return true
+	default:
+		return false
+	}
+}