From 52f8b557822e93d6976c9c10b4b8b6bb9233743d Mon Sep 17 00:00:00 2001 From: henry3260 Date: Tue, 22 Sep 2026 21:07:27 +0800 Subject: [PATCH 1/3] Reject a variadic task function when a handler is registered A variadic function passed the signature check as if its ... parameter were an ordinary slice, so it registered cleanly and then panicked inside reflect the first time the supervisor ran it: Execute calls the function with Call, which hands a variadic parameter the slice itself instead of spreading it. There is no argument shape the runtime can bind to a ... parameter, so the registration cannot become valid later. --- go-sdk/airflow/task_handler_test.go | 5 +++++ go-sdk/internal/bundle/task.go | 9 +++++++++ 2 files changed, 14 insertions(+) diff --git a/go-sdk/airflow/task_handler_test.go b/go-sdk/airflow/task_handler_test.go index 5df87af0e62c0..0112335935c56 100644 --- a/go-sdk/airflow/task_handler_test.go +++ b/go-sdk/airflow/task_handler_test.go @@ -120,6 +120,11 @@ func TestTaskHandlerPanicsOnBadHandler(t *testing.T) { fn: func(Context, chan int) error { return nil }, want: "type chan int cannot receive a task argument", }, + { + name: "variadic", + fn: func(Context, ...string) error { return nil }, + want: "is variadic; a task argument cannot fill a ... parameter", + }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/go-sdk/internal/bundle/task.go b/go-sdk/internal/bundle/task.go index c012b1a80aa14..c8febc84bee9c 100644 --- a/go-sdk/internal/bundle/task.go +++ b/go-sdk/internal/bundle/task.go @@ -156,6 +156,15 @@ func (f *taskFunction) validateFn(fnType reflect.Type) error { return fmt.Errorf("expected a func as input but was %s", fnType.Kind()) } + // Execute calls the function with Call, which passes a variadic parameter its slice + // rather than spreading it, so a variadic task function would panic at execution. + if fnType.IsVariadic() { + return fmt.Errorf( + "task function %s is variadic; a task argument cannot fill a ... parameter", + f.fullName, + ) + } + // Return values // `, error`, or just `error` if fnType.NumOut() < 1 || fnType.NumOut() > 2 { From 2e058b538e5e7b356e91ffc048dec8d0179a8e37 Mon Sep 17 00:00:00 2001 From: henry3260 Date: Tue, 22 Sep 2026 22:10:13 +0800 Subject: [PATCH 2/3] Close registration when a Go SDK bundle starts serving ADR 0007 states that registration closes when Serve is called, but nothing enforced it: a Register left below Serve in main silently added a handler to the map the runtime was already answering from, so what the bundle could run depended on how far main had got. Serve closes registration whatever the run goes on to do, since the mistake is where Register sits in main rather than what that run turned out to be. --- go-sdk/airflow/bundle.go | 19 ++++++++++++++++++- go-sdk/airflow/bundle_test.go | 22 ++++++++++++++++++++++ go-sdk/airflow/serve.go | 7 +++++++ 3 files changed, 47 insertions(+), 1 deletion(-) diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index b13b74ec76733..459878f6660d2 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -66,7 +66,8 @@ type Registerable interface{ registerable() } // // bundle.Register(reports.Handlers()...) // -// Register panics if a task handler with the same dag_id and task_id is already registered. +// Register panics if a task handler with the same dag_id and task_id is already registered, +// and if [BundleRef.Serve] has already been called: registration closes when serving starts. func (b *BundleRef) Register(items ...Registerable) { for _, item := range items { switch item := item.(type) { @@ -86,6 +87,7 @@ type taskHandlerMap struct { mu sync.RWMutex handlers map[string]map[string]bundle.Task order []bundle.TaskHandlerInfo + closed bool } var ( @@ -93,10 +95,25 @@ var ( _ bundle.EnumerableBundle = (*taskHandlerMap)(nil) ) +// close ends registration. Serve calls it before it starts answering for the bundle, so the +// map the runtime reads cannot gain a handler under it. +func (m *taskHandlerMap) close() { + m.mu.Lock() + defer m.mu.Unlock() + + m.closed = true +} + func (m *taskHandlerMap) add(dagId, taskId string, task bundle.Task) { m.mu.Lock() defer m.mu.Unlock() + if m.closed { + panic(fmt.Sprintf( + "airflow.BundleRef.Register: task %q of Dag %q was registered after Serve; "+ + "register every task handler before Serve", taskId, dagId, + )) + } if m.handlers == nil { m.handlers = make(map[string]map[string]bundle.Task) } diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index 9ad5217b30b6a..a0bebb0d6efc1 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -19,6 +19,7 @@ package airflow import ( "fmt" + "io" "os" "os/exec" "reflect" @@ -116,6 +117,27 @@ func TestRegisterRejectsDuplicateTask(t *testing.T) { ) } +func TestRegisterAfterServePanics(t *testing.T) { + b := Bundle() + b.Register(TaskHandler("py_etl", "transform", noop)) + require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard)) + + assert.PanicsWithValue(t, + `airflow.BundleRef.Register: task "load" of Dag "py_etl" was registered after Serve; `+ + `register every task handler before Serve`, + func() { b.Register(TaskHandler("py_etl", "load", noop)) }, + ) +} + +// Serve closes registration whatever the run does, so a bundle that only printed its usage +// still refuses a late Register. +func TestRegisterAfterAFailedServePanics(t *testing.T) { + b := Bundle() + require.NoError(t, b.serve([]string{"--help"}, io.Discard)) + + assert.Panics(t, func() { b.Register(TaskHandler("py_etl", "transform", noop)) }) +} + func TestRegisterIsSafeForConcurrentUse(t *testing.T) { const workers, perWorker = 8, 100 diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index eb61002a25e71..180404aaa2568 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -51,6 +51,9 @@ const ( // Serve runs the bundle. Call it as the last statement of main. // +// Serve closes registration: a [BundleRef.Register] call that reaches the bundle afterwards +// panics rather than changing what the running bundle answers for. +// // The command-line flags of the executable decide what Serve does. // With --airflow-metadata it prints the bundle's manifest and returns, which is how // airflow-go-pack reads the registered Dag and task ids. @@ -68,6 +71,10 @@ func (b *BundleRef) Serve() error { } func (b *BundleRef) serve(args []string, stdout io.Writer) error { + // Registration closes here, whatever this run turns out to do, so that a Register left + // below Serve in main is reported as the mistake it is rather than racing the runtime. + b.taskHandlers.close() + // The flags go on their own FlagSet. On pflag.CommandLine, every program that imports this // package would get them, and one that defines its own --format there would panic. flags := flag.NewFlagSet(os.Args[0], flag.ContinueOnError) From 3b1b028f9aa094c1a9911cbed81eff29ab07173a Mon Sep 17 00:00:00 2001 From: henry3260 Date: Wed, 23 Sep 2026 14:07:57 +0800 Subject: [PATCH 3/3] Close registration on the bundle rather than on its handler map MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The flag describes whether the bundle still takes registrations, not the state of one of the maps it holds. Keeping it on the handler map meant every further kind of registration would need a flag of its own and a matching line in Serve, so a native Dag would have arrived with the check missing until someone remembered it. The check now runs before Register looks at what the item is, so it covers a kind added later for free. It reads an atomic rather than taking a lock because the handler map keeps its own — the runtime is handed that map on its own — and a second mutex would introduce a lock order to respect. --- go-sdk/airflow/bundle.go | 25 +++++++++---------------- go-sdk/airflow/bundle_test.go | 17 ++++++++++++++--- go-sdk/airflow/serve.go | 2 +- 3 files changed, 24 insertions(+), 20 deletions(-) diff --git a/go-sdk/airflow/bundle.go b/go-sdk/airflow/bundle.go index 459878f6660d2..14522f47af984 100644 --- a/go-sdk/airflow/bundle.go +++ b/go-sdk/airflow/bundle.go @@ -21,6 +21,7 @@ import ( "fmt" "slices" "sync" + "sync/atomic" "github.com/apache/airflow/go-sdk/internal/bundle" ) @@ -28,6 +29,9 @@ import ( // BundleRef holds the task handlers that this executable runs for Airflow. // [Bundle] returns an empty one. type BundleRef struct { + // closed ends registration for everything the bundle can hold, so a kind added later + // is covered without a flag of its own. Serve sets it; Register reads it. + closed atomic.Bool taskHandlers taskHandlerMap } @@ -69,6 +73,11 @@ type Registerable interface{ registerable() } // Register panics if a task handler with the same dag_id and task_id is already registered, // and if [BundleRef.Serve] has already been called: registration closes when serving starts. func (b *BundleRef) Register(items ...Registerable) { + if b.closed.Load() { + panic( + "airflow.BundleRef.Register: Serve has already been called; register everything before Serve", + ) + } for _, item := range items { switch item := item.(type) { case *taskHandler: @@ -87,7 +96,6 @@ type taskHandlerMap struct { mu sync.RWMutex handlers map[string]map[string]bundle.Task order []bundle.TaskHandlerInfo - closed bool } var ( @@ -95,25 +103,10 @@ var ( _ bundle.EnumerableBundle = (*taskHandlerMap)(nil) ) -// close ends registration. Serve calls it before it starts answering for the bundle, so the -// map the runtime reads cannot gain a handler under it. -func (m *taskHandlerMap) close() { - m.mu.Lock() - defer m.mu.Unlock() - - m.closed = true -} - func (m *taskHandlerMap) add(dagId, taskId string, task bundle.Task) { m.mu.Lock() defer m.mu.Unlock() - if m.closed { - panic(fmt.Sprintf( - "airflow.BundleRef.Register: task %q of Dag %q was registered after Serve; "+ - "register every task handler before Serve", taskId, dagId, - )) - } if m.handlers == nil { m.handlers = make(map[string]map[string]bundle.Task) } diff --git a/go-sdk/airflow/bundle_test.go b/go-sdk/airflow/bundle_test.go index a0bebb0d6efc1..fb9503e9a3268 100644 --- a/go-sdk/airflow/bundle_test.go +++ b/go-sdk/airflow/bundle_test.go @@ -122,13 +122,24 @@ func TestRegisterAfterServePanics(t *testing.T) { b.Register(TaskHandler("py_etl", "transform", noop)) require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard)) - assert.PanicsWithValue(t, - `airflow.BundleRef.Register: task "load" of Dag "py_etl" was registered after Serve; `+ - `register every task handler before Serve`, + assert.PanicsWithValue( + t, + "airflow.BundleRef.Register: Serve has already been called; register everything before Serve", func() { b.Register(TaskHandler("py_etl", "load", noop)) }, ) } +// The flag lives on the bundle, not on the task-handler map, so a kind of item added to +// Register later is covered without a flag of its own. +func TestRegisterAfterServeRejectsEveryKindOfItem(t *testing.T) { + b := Bundle() + require.NoError(t, b.serve([]string{"--airflow-metadata"}, io.Discard)) + + var nilItem Registerable + assert.Panics(t, func() { b.Register(nilItem) }, + "the closed check runs before Register looks at what the item is") +} + // Serve closes registration whatever the run does, so a bundle that only printed its usage // still refuses a late Register. func TestRegisterAfterAFailedServePanics(t *testing.T) { diff --git a/go-sdk/airflow/serve.go b/go-sdk/airflow/serve.go index 180404aaa2568..f16cc6963f679 100644 --- a/go-sdk/airflow/serve.go +++ b/go-sdk/airflow/serve.go @@ -73,7 +73,7 @@ func (b *BundleRef) Serve() error { func (b *BundleRef) serve(args []string, stdout io.Writer) error { // Registration closes here, whatever this run turns out to do, so that a Register left // below Serve in main is reported as the mistake it is rather than racing the runtime. - b.taskHandlers.close() + b.closed.Store(true) // The flags go on their own FlagSet. On pflag.CommandLine, every program that imports this // package would get them, and one that defines its own --format there would panic.