[jonathansudhakar1]

Fixes #23550.

register.Combiner1/2/3 panicked with Failed to optimize MergeAccumulators for combiner &{}. Failed to infer types (and likewise for the other methods) when a CombineFn method takes a leading context.Context. CombineFn validation in graph already accepts that parameter, and it's how combiners use metrics and logging. Only register rejected it.

This follows the plan @lostluck agreed to on the issue (thank you!):

• register.tmpl: adds ctx-first variants of the combiner extraction interfaces (createAccumulator1x1/1x2, addInput3x1/3x2, mergeAccumulators3x1/3x2, extractOutput2x1/2x2). In each CombinerN, adds a branch per variant, next to its non-ctx counterpart. Each branch registers the caller for the func(context.Context, ...) type and builds the method wrapper, the same way the existing branches do.
• register.go: regenerated from the template. Only the combiner section changes.
• Compatibility: combiners without a context take the same branches as before. Signatures that are still unsupported still panic with "Failed to optimize".

Testing

All commands run from sdks/, Go 1.26.

• Regeneration. I built cmd/specialize and ran it the way doc.go's go:generate lines do (specialize --package=register --input=register.tmpl --x=data,universals --imports=typex, then gofmt). As a check, I first regenerated from the unchanged template: the output was byte-for-byte identical to the checked-in register.go. Then I regenerated from the edited template.
• New TestCombiner_ContextCombiners in register_test.go, table-driven, with 7 subtests:
  • Combiners covered: ctx in every method, with and without error returns, for Combiner1, Combiner2 and Combiner3. Together these cover all eight new interfaces, for each type pairing [T0, T0/T1/T2] that has an error-return variant. A seventh combiner, with only MergeAccumulators(ctx, a, b), is the exact repro from the issue.
  • Checks on the wrappers from reflectx.WrapMethods:
    • each wrapper has a leading context.Context parameter;
    • each uses the registered caller, not the reflection fallback;
    • each forwards the caller's ctx (the methods add a value read from ctx);
    • each returns the right results.
  • Checks on graph.NewCombineFn: it accepts the wrapped methods, and funcx detects their context parameter.
• go test ./go/pkg/beam/register/ -run TestCombiner_ContextCombiners -v: all 7 subtests PASS.
  With the pre-change register.go from upstream/master, the same test fails with the issue's panic: panic: Failed to optimize MergeAccumulators for combiner &{}. Failed to infer types.
• go test ./go/pkg/beam/register/...: ok. go vet ./go/pkg/beam/register/: clean. gofmt -l go/pkg/beam/register/: clean.
• go test ./go/pkg/beam/transforms/stats/... ./go/pkg/beam/transforms/top/... ./go/examples/snippets/... (other users of register.CombinerN): all ok.
• End to end on Prism (local only, not committed):
  • I ran beam.Combine with a registered combiner that takes a ctx in every method. Its AddInput increments a counter through that ctx.
  • It PASSES: the output is "55", and the counter is committed as 10.
  • Two mutants fail as expected: expecting "56" fails the passert, and expecting 11 for the counter fails the check.
  • On the old register.go it panics at init.

Pipeline test source

package register_test

import (
	"context"
	"strconv"
	"testing"

	"github.com/apache/beam/sdks/v2/go/pkg/beam"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/core/metrics"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/register"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/passert"
	"github.com/apache/beam/sdks/v2/go/pkg/beam/testing/ptest"
)

func TestMain(m *testing.M) {
	ptest.Main(m)
}

func init() {
	register.Combiner2[int, string](&sumWithContextFn{})
}

var sumWithContextInputs = beam.NewCounter("register_test", "sumWithContextInputs")

// sumWithContextFn sums ints and formats the result. Every method takes a
// context.Context, and AddInput uses it to update a counter.
type sumWithContextFn struct{}

func (fn *sumWithContextFn) CreateAccumulator(_ context.Context) int {
	return 0
}

func (fn *sumWithContextFn) AddInput(ctx context.Context, a, i int) (int, error) {
	sumWithContextInputs.Inc(ctx, 1)
	return a + i, nil
}

func (fn *sumWithContextFn) MergeAccumulators(_ context.Context, a, b int) int {
	return a + b
}

func (fn *sumWithContextFn) ExtractOutput(_ context.Context, a int) (string, error) {
	return strconv.Itoa(a), nil
}

func TestCombiner_ContextMethodsInPipeline(t *testing.T) {
	p, s, col := ptest.CreateList([]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10})
	sum := beam.Combine(s, &sumWithContextFn{}, col)
	passert.Equals(s, sum, "55")

	pr := ptest.RunAndValidate(t, p)

	qr := pr.Metrics().Query(func(sr metrics.SingleResult) bool {
		return sr.Name() == "sumWithContextInputs"
	})
	counters := qr.Counters()
	if len(counters) != 1 {
		t.Fatalf("pr.Metrics().Query(Name = \"sumWithContextInputs\") returned %d counters, want 1", len(counters))
	}
	if got, want := counters[0].Committed, int64(10); got != want {
		t.Errorf("sumWithContextInputs counter = %v, want %v", got, want)
	}
}

Open questions

• I left the Prism pipeline test out of the PR. In the register package's test binary, running any pipeline calls beam.Init. After that, the existing unit tests that call runtime.RegisterType inside test functions panic with "Init hooks have already run" (I tried it: TestEmitter1 fails). If you'd like an end-to-end test in the repo, I'm happy to add it somewhere better suited, for example sdks/go/test/integration/primitives. Let me know where you'd prefer it.

🤖 Generated with Claude Code

[github-actions]

📋 Review Summary

This PR adds support for optional leading context.Context parameters on CombineFn lifecycle methods (CreateAccumulator, AddInput, MergeAccumulators, and ExtractOutput) in register.Combiner1, Combiner2, and Combiner3, fixing #23550. The template changes in register.tmpl, regenerated code in register.go, and unit tests in register_test.go are clean, well-structured, and thoroughly tested.

🔍 General Feedback

• The new extraction interfaces and type-switch branches in register.tmpl accurately mirror the existing non-context patterns and align with graph.AsCombineFn signature validation.
• Unit test coverage in TestCombiner_ContextCombiners is comprehensive, verifying caller registration (avoiding reflection fallback), context propagation, error returns, and graph.NewCombineFn integration across all CombinerN variants.

---

💡 Tip: Comment /gemini review on this pull request to trigger a new Gemini code review.

[github-actions]

Assigning reviewers:

R: @lostluck for label go.
Backup reviewer: shunping (if the main reviewer is unresponsive, you can escalate to them by tagging their username)

Note: If you would like to opt out of this review, comment assign to next reviewer.

Available commands:

• stop reviewer notifications - opt out of the automated review tooling
• remind me after tests pass - tag the comment author after tests pass
• waiting on author - shift the attention set back to the author (any comment or push by the author will return the attention set to the reviewers)

The PR bot will only process comments in the main thread (not review comments).