[Go SDK] wrap a failed last data/timer send like Flush

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)).

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 {

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.