diff --git a/internal/cli/examples_test.go b/internal/cli/examples_test.go index b3ea7ed1..3c502de0 100644 --- a/internal/cli/examples_test.go +++ b/internal/cli/examples_test.go @@ -6,6 +6,7 @@ import ( "os" "path/filepath" "strings" + "sync" "testing" "time" @@ -166,7 +167,10 @@ func TestConfigValidation_ExampleConfigsBuildRealComponents(t *testing.T) { // nothing the commands above may have failed to create. Building // it second meant a handler that skipped for a missing ATTACHed // table took the sink check down with it. - _, err = sinks.New(ctx, conf.Pipeline.Sink, conn) + // The run command hands every sink the connection's lock, and a + // sqlcommand sink refuses to build without one. + lock := &sync.Mutex{} + _, err = sinks.New(ctx, conf.Pipeline.Sink, conn, sinks.WithConnLock(lock)) checkBuildError(t, "sink", err) if conf.Tables != nil { @@ -174,7 +178,7 @@ func TestConfigValidation_ExampleConfigsBuildRealComponents(t *testing.T) { if table.Window == nil { continue } - _, err := sinks.New(ctx, table.Window.Sink, conn) + _, err := sinks.New(ctx, table.Window.Sink, conn, sinks.WithConnLock(lock)) checkBuildError(t, "window sink for "+table.Name, err) } } diff --git a/internal/cli/run/managers.go b/internal/cli/run/managers.go index 378d9957..d1d7e64d 100644 --- a/internal/cli/run/managers.go +++ b/internal/cli/run/managers.go @@ -110,10 +110,16 @@ func buildManagedTables( "table %q: opening the window sink's connection", table.Name) } conns = append(conns, sinkConn) + // The sink's connection is the window's alone, and the manager runs + // the sink from one goroutine, so the lock a sqlcommand sink needs + // has no other party. It is a lock of its own rather than the + // pipeline's: taking the pipeline's here would stall the consume + // loop for the length of every window write. sink, err := sinks.New(ctx, table.Window.Sink, sinkConn, sinks.WithMeterProvider(mp), sinks.WithSinkRole("manager"), - sinks.WithRetryEvents(events)) + sinks.WithRetryEvents(events), + sinks.WithConnLock(&sync.Mutex{})) if err != nil { return nil, closeConns, fmt.Errorf("table %q window sink: %w", table.Name, err) } diff --git a/internal/cli/run/root.go b/internal/cli/run/root.go index d018386b..7381917e 100644 --- a/internal/cli/run/root.go +++ b/internal/cli/run/root.go @@ -37,6 +37,7 @@ func newErrorPolicies( ctx context.Context, conf *config.Conf, conn adbc.Connection, + lock *sync.Mutex, mp metric.MeterProvider, events sinks.RetryEvents, ) (core.PipelineErrorPolicies, error) { @@ -63,7 +64,8 @@ func newErrorPolicies( dlqSink, err := sinks.New(ctx, *onError.DLQ, conn, sinks.WithMeterProvider(mp), sinks.WithSinkRole("dlq"), - sinks.WithRetryEvents(events)) + sinks.WithRetryEvents(events), + sinks.WithConnLock(lock)) if err != nil { return policies, fmt.Errorf("pipeline.on_error dlq: %w", err) } @@ -73,6 +75,23 @@ func newErrorPolicies( return policies, nil } +// newPipelineSink builds the pipeline's own sink. A function of its own so the +// shared-connection test builds it exactly as the run command does. +func newPipelineSink( + ctx context.Context, + conf *config.Conf, + conn adbc.Connection, + lock *sync.Mutex, + mp metric.MeterProvider, + events sinks.RetryEvents, +) (core.Sink, error) { + return sinks.New(ctx, conf.Pipeline.Sink, conn, + sinks.WithMeterProvider(mp), + sinks.WithSinkRole(core.SinkRolePipeline), + sinks.WithRetryEvents(events), + sinks.WithConnLock(lock)) +} + func NewCommand() *cobra.Command { var configPath string var maxMsgs int @@ -375,10 +394,7 @@ func NewCommand() *cobra.Command { // The signal context, so a SIGTERM arriving while the sink dials // its destination stops the start instead of waiting it out. - sink, err := sinks.New(ctx, conf.Pipeline.Sink, conn, - sinks.WithMeterProvider(meterProvider), - sinks.WithSinkRole(core.SinkRolePipeline), - sinks.WithRetryEvents(retryEvents)) + sink, err := newPipelineSink(ctx, conf, conn, lock, meterProvider, retryEvents) if err != nil { return err } @@ -401,7 +417,7 @@ func NewCommand() *cobra.Command { }() } - errorPolicies, err := newErrorPolicies(ctx, conf, conn, meterProvider, retryEvents) + errorPolicies, err := newErrorPolicies(ctx, conf, conn, lock, meterProvider, retryEvents) if err != nil { return err } diff --git a/internal/cli/run/row_roles_test.go b/internal/cli/run/row_roles_test.go index 51c6606b..9e57e191 100644 --- a/internal/cli/run/row_roles_test.go +++ b/internal/cli/run/row_roles_test.go @@ -3,6 +3,7 @@ package run import ( "context" "os" + "sync" "testing" "github.com/apache/arrow-adbc/go/adbc" @@ -112,7 +113,7 @@ func TestDLQRowsCarryTheDLQRole(t *testing.T) { DLQ: &config.Sink{Type: "console"}, } - policies, err := newErrorPolicies(context.Background(), conf, nil, mp, sinks.RetryEvents{}) + policies, err := newErrorPolicies(context.Background(), conf, nil, &sync.Mutex{}, mp, sinks.RetryEvents{}) assert.NoError(t, err) assert.That(t, policies.DLQSink != nil) diff --git a/internal/cli/run/shared_conn_test.go b/internal/cli/run/shared_conn_test.go new file mode 100644 index 00000000..6a0f13fa --- /dev/null +++ b/internal/cli/run/shared_conn_test.go @@ -0,0 +1,168 @@ +package run + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow/array" + "github.com/turbolytics/sql-flow/internal/config" + "github.com/turbolytics/sql-flow/internal/core" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/turbolytics/sql-flow/internal/handlers" + "github.com/turbolytics/sql-flow/internal/sinks" + "github.com/zeebo/assert" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.uber.org/zap" +) + +// sliceSource delivers fixed batches and then closes its stream. +type sliceSource struct { + batches [][]core.Message +} + +func (s *sliceSource) Start() error { return nil } +func (s *sliceSource) Commit() error { return nil } +func (s *sliceSource) Close() error { return nil } + +func (s *sliceSource) Stream() <-chan []core.Message { + ch := make(chan []core.Message, len(s.batches)) + for _, b := range s.batches { + ch <- b + } + close(ch) + return ch +} + +func sharedConnCount(t *testing.T, conn adbc.Connection, table string) int64 { + t.Helper() + stmt, err := conn.NewStatement() + assert.NoError(t, err) + defer stmt.Close() + assert.NoError(t, stmt.SetSqlQuery("SELECT count(*) FROM "+table)) + reader, _, err := stmt.ExecuteQuery(context.Background()) + assert.NoError(t, err) + defer reader.Release() + assert.That(t, reader.Next()) + return reader.Record().Column(0).(*array.Int64).Value(0) +} + +// Every party that runs statements on the pipeline's DuckDB connection holds +// the shared lock while it does: the handler, the pipeline sink, the DLQ sink, +// the progress store, and each table manager with its own sink. +// +// DuckDB closes a pending result the moment another statement runs on its +// connection. A party that skips the lock fails whichever query another party +// has in flight, with "Attempting to execute an unsuccessful or closed pending +// query result". #280: the handler reset and the progress write ran outside +// the lock, the Bluesky demo failed a window poll at batch 500, and at batch 1 +// the failure recurred every few minutes. On main a failed poll stops the +// process. +// +// The race has no deterministic reproduction, so this checks the invariant on +// every statement instead. The components are built by the same functions the +// run command calls, over a connection that records any execution finding the +// lock free. Nothing else touches the connection while the loop and the poll +// run, so a free lock means a party let go of it. +// +// It covers the two places that build a sqlcommand sink on the shared +// connection, the pipeline's and the DLQ's. sinks.New refuses to build one +// without the lock; this test is what shows the sink then uses it. The window +// manager is built beside them and polls, but it runs on connections of its +// own since #281, so it is not a party to the shared one: the lock checker +// never sees it, and that is the point. +func TestPipelineSharedConnection_EveryPartyHoldsTheLock(t *testing.T) { + coverage.Covers(t, "manager.window", "sink.sqlcommand", "error.dlq", "handler.inferred_mem") + ctx := context.Background() + + db, raw := rowsTestDB(t) + rowsTestExec(t, raw, "CREATE TABLE win (bucket TIMESTAMPTZ, n BIGINT)") + rowsTestExec(t, raw, "CREATE TABLE published (bucket TIMESTAMPTZ, n BIGINT)") + rowsTestExec(t, raw, "CREATE TABLE dead (error VARCHAR, message VARCHAR, phase VARCHAR, timestamp VARCHAR)") + + lock := &sync.Mutex{} + conn := duckdb.NewLockChecked(raw, lock) + mp := sdkmetric.NewMeterProvider() + + conf := &config.Conf{ + Tables: &config.Tables{ + SQL: []config.TableSQL{{ + Name: "win", + // One-second buckets with no grace: a bucket closes once a + // later one exists. + Window: &config.Window{ + TimeColumn: "bucket", + SizeSeconds: 1, + LateRows: "drop", + PollIntervalSecs: 3600, + EmitSQL: "SELECT bucket, sum(n)::BIGINT AS n FROM closed GROUP BY ALL", + Sink: config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{ + SQL: "INSERT INTO published SELECT bucket, n FROM sqlflow_sink_batch", + }}, + }, + }}, + }, + } + conf.Pipeline.Handler = config.Handler{ + Type: "handlers.InferredMemBatch", + SQL: "SELECT to_timestamp(a) AS bucket, count(*)::BIGINT AS n FROM batch GROUP BY ALL", + } + conf.Pipeline.Sink = config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{ + SQL: "INSERT INTO win SELECT bucket, n FROM sqlflow_sink_batch", + }} + conf.Pipeline.OnError = &config.OnError{ + Policy: "DLQ", + DLQ: &config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{ + SQL: "INSERT INTO dead SELECT * FROM sqlflow_sink_batch", + }}, + } + + // Startup runs before any second goroutine exists, and the run command + // does it without the lock. Holding it here keeps the check about the + // running pipeline. + progress := core.NewProgressStore(conn) + lock.Lock() + assert.NoError(t, progress.Init(ctx)) + assert.NoError(t, initWindowStores(ctx, conf, conn)) + lock.Unlock() + + sink, err := newPipelineSink(ctx, conf, conn, lock, mp, sinks.RetryEvents{}) + assert.NoError(t, err) + policies, err := newErrorPolicies(ctx, conf, conn, lock, mp, sinks.RetryEvents{}) + assert.NoError(t, err) + handler, err := handlers.New(conn, conf.Pipeline.Handler, zap.NewNop()) + assert.NoError(t, err) + managed, closeConns, err := buildManagedTables(ctx, conf, db, zap.NewNop(), mp, + nil, sinks.RetryEvents{}) + assert.NoError(t, err) + defer closeConns() + assert.Equal(t, 1, len(managed)) + + msg := func(v string) core.Message { return core.Message{Value: []byte(v)} } + src := &sliceSource{batches: [][]core.Message{ + {msg(`{"a": 1}`), msg(`{"a": 1}`)}, + {msg(`not json`), msg(`{"a": 2}`)}, + {msg(`{"a": 2}`), msg(`{"a": 3}`)}, + }} + tb := core.NewTurbine(src, handler, sink, 2, time.Hour, lock, policies, + core.WithProgressStore(progress)) + + _, err = tb.ConsumeLoop(ctx, 6) + assert.NoError(t, err) + assert.NoError(t, managed[0].Poll(ctx)) + + // Every party ran at least once, or a clean result proves nothing. The + // window closed the two buckets a later one exists for, and the newest + // stays open. + assert.Equal(t, int64(2), sharedConnCount(t, raw, "published")) + assert.Equal(t, int64(1), sharedConnCount(t, raw, "dead")) + + if v := conn.Violations(); len(v) > 0 { + t.Fatalf("%d statements ran on the shared connection without the lock:\n%s", + len(v), strings.Join(v, "\n")) + } +} diff --git a/internal/core/lock_test.go b/internal/core/lock_test.go new file mode 100644 index 00000000..8c2cce4a --- /dev/null +++ b/internal/core/lock_test.go @@ -0,0 +1,96 @@ +package core + +import ( + "context" + "strings" + "sync" + "testing" + "time" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow" + "github.com/turbolytics/sql-flow/internal/coverage" + "github.com/turbolytics/sql-flow/internal/duckdb" + "github.com/zeebo/assert" +) + +// connHandler is a fakeHandler that touches the connection the way the real +// handlers do: Init resets a batch table and Invoke queries one. The pipeline +// decides whether either runs under the lock, and the fake handler that +// touches nothing hides that decision. +type connHandler struct { + fakeHandler + conn adbc.Connection +} + +func (h *connHandler) exec(ctx context.Context, sql string) error { + stmt, err := h.conn.NewStatement() + if err != nil { + return err + } + defer stmt.Close() + if err := stmt.SetSqlQuery(sql); err != nil { + return err + } + _, err = stmt.ExecuteUpdate(ctx) + return err +} + +func (h *connHandler) Init(ctx context.Context) error { + if err := h.fakeHandler.Init(ctx); err != nil { + return err + } + return h.exec(ctx, "DROP TABLE IF EXISTS batch") +} + +func (h *connHandler) Invoke(ctx context.Context) (arrow.Table, error) { + if err := h.exec(ctx, "CREATE OR REPLACE TABLE batch AS SELECT 1 AS n"); err != nil { + return nil, err + } + return h.fakeHandler.Invoke(ctx) +} + +// Every statement the pipeline runs on the shared connection runs under the +// shared lock. The table managers collect on the same connection under that +// lock, and DuckDB closes a pending result the moment another statement runs +// on its connection, so a statement outside the lock fails whichever poll is +// in flight. #280: the handler reset after every batch and the progress write +// both ran outside the lock, and the Bluesky demo logged "closed pending +// query result" once at batch 500 and seven times in ten minutes at batch 1. +// +// The race cannot be reproduced on demand. The invariant can be checked on +// every statement: the wrapper records any execution that finds the lock +// free. The loop runs with no other goroutine touching the connection, so a +// free lock means the pipeline itself let go of it. +func TestCoreConsumeLoop_EveryStatementOnTheSharedConnectionHoldsTheLock(t *testing.T) { + coverage.Covers(t, "state.durability") + ctx := context.Background() + db, err := duckdb.OpenPath(ctx, "") + assert.NoError(t, err) + defer db.Close() + raw, err := db.Connect(ctx) + assert.NoError(t, err) + defer raw.Close() + + lock := &sync.Mutex{} + conn := duckdb.NewLockChecked(raw, lock) + + // Startup runs before the managers exist, and root.go does it without + // the lock. Holding it here keeps the check about the running pipeline. + progress := NewProgressStore(conn) + lock.Lock() + assert.NoError(t, progress.Init(ctx)) + lock.Unlock() + + src := &fakeSource{batches: [][]Message{messages(4), messages(4), messages(4)}} + tb := NewTurbine(src, &connHandler{conn: conn}, &fakeSink{}, 4, time.Second, + lock, PipelineErrorPolicies{}, WithProgressStore(progress)) + stats, err := tb.ConsumeLoop(ctx, 12) + assert.NoError(t, err) + assert.Equal(t, int64(12), stats.MessagesConsumed()) + + if v := conn.Violations(); len(v) > 0 { + t.Fatalf("%d statements ran on the shared connection without the lock:\n%s", + len(v), strings.Join(v, "\n")) + } +} diff --git a/internal/core/turbine.go b/internal/core/turbine.go index a16f22cd..fbf60f63 100644 --- a/internal/core/turbine.go +++ b/internal/core/turbine.go @@ -388,7 +388,13 @@ func (t *Turbine) recordProgress(ctx context.Context) { return } now := time.Now().UTC() + // Held through the write below. The table managers collect on this same + // connection under this lock, and DuckDB closes a pending result the + // moment another statement runs on its connection. A write outside the + // lock failed whichever poll was in flight with "closed pending query + // result", at a rate set by batch size (#280). t.lock.Lock() + defer t.lock.Unlock() p := Progress{LastCommit: now, Messages: t.stats.MessagesConsumed()} if !t.arrivedAt.IsZero() { p.LastArrival = t.arrivedAt @@ -413,7 +419,6 @@ func (t *Turbine) recordProgress(ctx context.Context) { t.progressWrittenAt = now t.writtenArrival = t.arrivedAt } - t.lock.Unlock() // The snapshot above is exact and free. The table is not: one UPDATE // through ADBC measures about 112 microseconds, and a batch of 5000 at a @@ -594,7 +599,8 @@ func (t *Turbine) ConsumeLoop(ctx context.Context, maxMsgs int) (stats *Stats, e t.stats.StartTime = time.Now().UTC() t.stats.SetNumMessagesConsumed(0) - if err := t.handler.Init(ctx); err != nil { + // Under the lock: the managers are already polling this connection (#280). + if err := t.initHandler(ctx); err != nil { return nil, err } @@ -1033,6 +1039,18 @@ func (t *Turbine) commitSource() error { return t.source.Commit() } +// initHandler resets the handler under the lock. Init drops or truncates the +// batch table on the shared connection, and the table managers poll that +// connection from their own goroutines. DuckDB closes a pending result the +// moment another statement runs on its connection, so an unlocked reset +// failed whichever collect was in flight with "closed pending query result", +// once per batch at batch size 1 (#280). +func (t *Turbine) initHandler(ctx context.Context) error { + t.lock.Lock() + defer t.lock.Unlock() + return t.handler.Init(ctx) +} + // rollbackState discards this batch's uncommitted state writes. Used on the // paths that fail before commitState is reached. func (t *Turbine) rollbackState(ctx context.Context) { @@ -1272,7 +1290,7 @@ func (t *Turbine) processBatch(ctx context.Context, numBatchMessages int) error } i0 := time.Now() - err = t.handler.Init(ctx) + err = t.initHandler(ctx) t.recordPhase(ctx, phaseHandlerInit, time.Since(i0)) if err != nil { diff --git a/internal/duckdb/lockcheck.go b/internal/duckdb/lockcheck.go new file mode 100644 index 00000000..7f7c9557 --- /dev/null +++ b/internal/duckdb/lockcheck.go @@ -0,0 +1,157 @@ +package duckdb + +import ( + "context" + "fmt" + "runtime" + "strings" + "sync" + + "github.com/apache/arrow-adbc/go/adbc" + "github.com/apache/arrow-go/v18/arrow/array" +) + +// LockChecked wraps a connection so that every statement executed on it, and +// every commit or rollback, must run while the given lock is held. It is a +// test aid for the one rule the pipeline has about its DuckDB connection. +// +// The pipeline, the table managers, the progress store and the debug API all +// share one connection, serialized by one mutex. DuckDB closes a pending +// result on a connection the moment another statement runs on it, so a +// statement that skips the lock does not deadlock or corrupt anything: it +// fails whatever query was in flight, with "Attempting to execute an +// unsuccessful or closed pending query result". The failure lands on the +// other party, at a rate set by traffic, which is why #280 shipped and ran +// for two days before a batch size of one made it reproducible. +// +// A race like that has no deterministic reproduction. The invariant does: +// at the moment a statement executes, the lock is held. TryLock succeeding +// means nobody held it, which is a violation whoever the caller is. TryLock +// failing means someone held it, and this wrapper cannot tell whether that +// someone is the caller; a concurrent holder hides the violation for that +// one call. So run the pipeline single-threaded under the wrapper, and every +// unlocked statement is caught on its first execution. +type LockChecked struct { + adbc.Connection + lock *sync.Mutex + + mu sync.Mutex + violations []string +} + +// NewLockChecked returns conn wrapped so that executing on it without lock +// held is recorded. Wrap after startup: the tables are created before any +// second goroutine exists, and root.go creates them without the lock. +func NewLockChecked(conn adbc.Connection, lock *sync.Mutex) *LockChecked { + return &LockChecked{Connection: conn, lock: lock} +} + +// Violations lists every execution that ran without the lock, each with the +// operation and the frames that led to it. +func (c *LockChecked) Violations() []string { + c.mu.Lock() + defer c.mu.Unlock() + return append([]string(nil), c.violations...) +} + +func (c *LockChecked) check(op string) { + if !c.lock.TryLock() { + return + } + c.lock.Unlock() + + pcs := make([]uintptr, 12) + n := runtime.Callers(3, pcs) + frames := runtime.CallersFrames(pcs[:n]) + var where []string + for { + f, more := frames.Next() + if !strings.HasSuffix(f.File, "lockcheck.go") { + where = append(where, fmt.Sprintf("%s:%d", f.Function, f.Line)) + } + if !more || len(where) == 6 { + break + } + } + + c.mu.Lock() + c.violations = append(c.violations, op+" without the lock\n\t"+strings.Join(where, "\n\t")) + c.mu.Unlock() +} + +func (c *LockChecked) Commit(ctx context.Context) error { + c.check("Commit") + return c.Connection.Commit(ctx) +} + +func (c *LockChecked) Rollback(ctx context.Context) error { + c.check("Rollback") + return c.Connection.Rollback(ctx) +} + +func (c *LockChecked) NewStatement() (adbc.Statement, error) { + stmt, err := c.Connection.NewStatement() + if err != nil { + return nil, err + } + return &lockCheckedStatement{Statement: stmt, conn: c}, nil +} + +// SetOption forwards to the wrapped connection's post-init options, which +// root.go uses to turn autocommit off for a state path. +func (c *LockChecked) SetOption(key, val string) error { + po, ok := c.Connection.(adbc.PostInitOptions) + if !ok { + return fmt.Errorf("lockcheck: wrapped connection does not support SetOption") + } + return po.SetOption(key, val) +} + +type lockCheckedStatement struct { + adbc.Statement + conn *LockChecked + sql string +} + +func (s *lockCheckedStatement) SetSqlQuery(query string) error { + s.sql = query + return s.Statement.SetSqlQuery(query) +} + +func (s *lockCheckedStatement) ExecuteQuery(ctx context.Context) (array.RecordReader, int64, error) { + s.conn.check("ExecuteQuery " + firstLine(s.sql)) + reader, n, err := s.Statement.ExecuteQuery(ctx) + if err != nil { + return nil, n, err + } + return &lockCheckedReader{RecordReader: reader, stmt: s}, n, nil +} + +func (s *lockCheckedStatement) ExecuteUpdate(ctx context.Context) (int64, error) { + s.conn.check("ExecuteUpdate " + firstLine(s.sql)) + return s.Statement.ExecuteUpdate(ctx) +} + +// lockCheckedReader checks each Next. A result is streamed, so a reader +// drained after the lock is released is as exposed as a statement executed +// without it. +type lockCheckedReader struct { + array.RecordReader + stmt *lockCheckedStatement +} + +func (r *lockCheckedReader) Next() bool { + r.stmt.conn.check("Next " + firstLine(r.stmt.sql)) + return r.RecordReader.Next() +} + +func firstLine(sql string) string { + sql = strings.TrimSpace(sql) + if i := strings.IndexByte(sql, '\n'); i >= 0 { + sql = sql[:i] + } + if len(sql) > 60 { + sql = sql[:60] + } + return sql +} diff --git a/internal/sinks/init.go b/internal/sinks/init.go index 5dedef6b..0bf95869 100644 --- a/internal/sinks/init.go +++ b/internal/sinks/init.go @@ -3,6 +3,7 @@ package sinks import ( "context" "sort" + "sync" "github.com/apache/arrow-adbc/go/adbc" "github.com/apache/arrow-go/v18/arrow" @@ -31,6 +32,14 @@ type options struct { meterProvider metric.MeterProvider role string retryEvents RetryEvents + connLock *sync.Mutex +} + +// WithConnLock supplies the lock that serializes the pipeline's DuckDB +// connection. The sqlcommand sink runs its statements on that connection, so +// New refuses to build one without it. See SQLCommandSink.connLock. +func WithConnLock(lock *sync.Mutex) Option { + return func(o *options) { o.connLock = lock } } // RetryEvents is told when a sink's retry ladder runs. Retry fires per failed @@ -94,6 +103,18 @@ func New(ctx context.Context, sink config.Sink, conn adbc.Connection, opts ...Op return nil, err } + // Checked here rather than in the builder, so every caller of New is held + // to it: the pipeline, each table manager and the DLQ build their sinks in + // three different places, and one of them forgetting the lock is a race + // that no single-threaded test of that sink would see. + if s, ok := built.(*SQLCommandSink); ok { + if o.connLock == nil { + return nil, errs.New(errs.CodeSinkInvalid, + "sink: sqlcommand shares the pipeline's DuckDB connection and needs its lock") + } + s.connLock = o.connLock + } + // Checked before the pipeline consumes anything, so a destination that is // not there fails the start rather than the first flush. if err := probe(ctx, built); err != nil { diff --git a/internal/sinks/sqlcommand.go b/internal/sinks/sqlcommand.go index ef97945d..d0ab2878 100644 --- a/internal/sinks/sqlcommand.go +++ b/internal/sinks/sqlcommand.go @@ -26,6 +26,12 @@ type SQLCommandSink struct { sql string substitutions []config.SQLCommandSubstitution + // connLock serializes conn with the pipeline, the table managers and the + // debug API, which all run statements on it. New sets it. Nil leaves + // serialization to the caller, which only a test driving the sink alone + // can promise. + connLock *sync.Mutex + mu sync.Mutex tables []arrow.Table } @@ -93,6 +99,17 @@ func (s *SQLCommandSink) requeue(tables []arrow.Table) { // send is one delivery attempt. It releases nothing: whether these batches can // be dropped is the caller's decision, and it depends on this error. func (s *SQLCommandSink) send(ctx context.Context, tables []arrow.Table) error { + // Held across the drop, the ingest and the user's SQL. The handler and the + // table managers run statements on this connection from other goroutines, + // and DuckDB closes a pending result the moment another statement runs on + // its connection, so an unlocked flush failed whatever query they had in + // flight (#280). Nobody calls Flush holding this lock: the pipeline and the + // managers release it before they flush. + if s.connLock != nil { + s.connLock.Lock() + defer s.connLock.Unlock() + } + if err := s.materialize(ctx, tables); err != nil { return err } diff --git a/internal/sinks/sqlcommand_test.go b/internal/sinks/sqlcommand_test.go index f27a8c0a..bd98593b 100644 --- a/internal/sinks/sqlcommand_test.go +++ b/internal/sinks/sqlcommand_test.go @@ -8,6 +8,7 @@ package sinks import ( "context" "strings" + "sync" "testing" "github.com/apache/arrow-adbc/go/adbc" @@ -109,6 +110,23 @@ func TestSinkSqlcommand_NewRequiresSQL(t *testing.T) { assert.Error(t, err) } +// The sink runs statements on the pipeline's shared connection, so New refuses +// to build one without the lock that serializes it. A call site that forgets +// the option then fails at startup, not with a "closed pending query result" +// in some other party's query at a rate set by traffic (#280). +func TestSinkSqlcommand_NewRefusesWithoutTheConnectionLock(t *testing.T) { + coverage.Covers(t, "sink.sqlcommand") + conn := newSinkTestConn(t) + conf := config.Sink{Type: "sqlcommand", SQLCommand: &config.SQLCommandSink{SQL: "SELECT 1"}} + + _, err := New(context.Background(), conf, conn) + assert.Error(t, err) + assert.That(t, strings.Contains(err.Error(), "needs its lock")) + + _, err = New(context.Background(), conf, conn, WithConnLock(&sync.Mutex{})) + assert.NoError(t, err) +} + // The batch reaches the user's SQL under the name the Python sink registered // it as. A different name and every shipped sqlcommand config stops working. func TestSinkSqlcommand_RunsSQLAgainstTheBatch(t *testing.T) {