From 2221917229d948f5d73caabb743d8fc7c23ae2f6 Mon Sep 17 00:00:00 2001 From: Igor Date: Thu, 1 Oct 2026 12:06:29 +0300 Subject: [PATCH 1/2] [Go SDK] Attribute a failed last data/timer send like Flush --- CHANGES.md | 1 + .../pkg/beam/core/runtime/harness/datamgr.go | 10 ++++- .../beam/core/runtime/harness/datamgr_test.go | 38 +++++++++++++++++++ 3 files changed, 47 insertions(+), 2 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index dbe5a0544556..87164fb5b9f4 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -84,6 +84,7 @@ * (Go) Fixed a data race on the Prism runner's artifact cache map in JobServices ([#32656](https://github.com/apache/beam/issues/32656)). * (Go) Fixed the harness leaking Data/State gRPC streams after the worker stops, and a deadlock when Send returns EOF ([#40260](https://github.com/apache/beam/issues/40260)). +* (Go) Wrap last dataWriter/timerWriter send errors on Close ([#40405](https://github.com/apache/beam/issues/40405)). * (Java) Fixed the declared schema of the error output of the Kafka write SchemaTransform, which wrapped the error schema a second time and did not match the rows it emits ([#39760](https://github.com/apache/beam/issues/39760)). * (Go) Fixed pubsubio importing a `google.golang.org/genproto` package removed in recent releases, which broke builds of Go modules depending on a current `genproto` version ([#40018](https://github.com/apache/beam/issues/40018)). * (Java) BigQueryIO now treats a 404 when deleting a temporary table or dataset as success, so a replayed work item whose earlier attempt already deleted it no longer retries forever ([#24997](https://github.com/apache/beam/issues/24997)). diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go index b9718e9b42b9..9905d9dc3efe 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go +++ b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go @@ -615,7 +615,10 @@ func (w *dataWriter) Close() error { }, }, } - return w.send(msg) + if err := w.send(msg); err != nil { + return errors.Wrapf(err, "dataWriter[%v;%v].Close: sending last data", w.id, w.ch.id) + } + return nil } const largeBufferNotificationThreshold = 1024 * 1024 * 1024 // 1GB @@ -723,7 +726,10 @@ func (w *timerWriter) Close() error { }, }, } - return w.send(msg) + if err := w.send(msg); err != nil { + return errors.Wrapf(err, "timerWriter[%v;%v].Close: sending last timers", w.id, w.ch.id) + } + return nil } func (w *timerWriter) writeTimers(p []byte) error { diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go b/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go index 637076d6e44c..5cd90dc80857 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go +++ b/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go @@ -18,6 +18,7 @@ package harness import ( "bytes" "context" + "errors" "fmt" "io" "log" @@ -644,6 +645,43 @@ func TestTimerWriterSendEOF(t *testing.T) { } } +type failLastClient struct{ block chan struct{} } + +func (f *failLastClient) Recv() (*fnpb.Elements, error) { <-f.block; return nil, io.EOF } +func (f *failLastClient) Send(e *fnpb.Elements) error { + for _, d := range e.GetData() { + if d.GetIsLast() { + return errors.New("Unavailable") + } + } + for _, tm := range e.GetTimers() { + if tm.GetIsLast() { + return errors.New("Unavailable") + } + } + return nil +} + +func TestDataWriter_Close(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + client := &failLastClient{block: make(chan struct{})} + defer close(client.block) + c := makeDataChannel(ctx, "id", client, cancel) + + w := c.OpenWrite(ctx, "ptr", "inst1") + if _, err := w.Write([]byte("element")); err != nil { + t.Fatal(err) + } + if err := w.Close(); err == nil || !strings.Contains(err.Error(), "dataWriter[") { + t.Fatalf("Close = %v, want dataWriter attribution", err) + } + tw := c.OpenTimerWrite(ctx, "ptr", "inst2", "family") + if err := tw.Close(); err == nil || !strings.Contains(err.Error(), "timerWriter[") { + t.Fatalf("timer Close = %v, want timerWriter attribution", err) + } +} + type noopDataClient struct { } From 4ac2e09c39c6577e0433da4535b299d0f3bda8de Mon Sep 17 00:00:00 2001 From: Igor Date: Sun, 11 Oct 2026 15:22:23 +0300 Subject: [PATCH 2/2] Fix timer writer error handling in tests --- sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go b/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go index 4a8b69fe7cd2..359758dc1cac 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go +++ b/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go @@ -682,6 +682,8 @@ func TestDataWriter_Close(t *testing.T) { if err := tw.Close(); err == nil || !strings.Contains(err.Error(), "timerWriter[") { t.Fatalf("timer Close = %v, want timerWriter attribution", err) } +} + // Production Open against closeInstruction. // Flush holds ch.mu inside client.Send. // Close takes m.mu and waits on ch.mu.