diff --git a/CHANGES.md b/CHANGES.md index cb5eced42f7c..d98c070c54a4 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -91,6 +91,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)). * (Go) Fixed a deadlock recreating a data channel while holding the channel lock ([#40414](https://github.com/apache/beam/issues/40414)). * (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)). diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go index 66ebd5aa6f9f..6831173cea1d 100644 --- a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go +++ b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go @@ -618,7 +618,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 @@ -726,7 +729,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 b48967daf9c5..359758dc1cac 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" @@ -646,6 +647,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) + } +} + // Production Open against closeInstruction. // Flush holds ch.mu inside client.Send. // Close takes m.mu and waits on ch.mu.