Close harness data/state channels and don't Recv on Send EOF (#40262)

* [Go SDK] Close data and state channels on harness exit

* [Go SDK] Don't Recv on write after Send EOF

* [Go SDK] Treat canceled StateChannel reads as closed

* Document bugfix for Go harness issues

Fixed harness leaking Data/State gRPC streams and a deadlock issue.

* Fix empty critical section in tests

---------

Co-authored-by: Robert Burke <lostluck@users.noreply.github.com>

Author: igor-dv

CHANGES.md

@@ -82,6 +82,7 @@
 ## Bugfixes
 
 * (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)).
 * (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)).

sdks/go/pkg/beam/core/runtime/harness/datamgr.go

@@ -144,6 +144,20 @@ func (m *DataChannelManager) Open(ctx context.Context, port exec.Port) (*DataCha
 	return ch, nil
 }
 
+// Close closes all cached DataChannels.
+func (m *DataChannelManager) Close() {
+	m.mu.Lock()
+	chans := m.ports
+	m.ports = nil
+	m.mu.Unlock()
+	for _, ch := range chans {
+		ch.mu.Lock()
+		ch.forceRecreate = nil
+		ch.mu.Unlock()
+		ch.cancelFn()
+	}
+}
+
 func (m *DataChannelManager) closeInstruction(instID instructionID, ports []exec.Port) error {
 	m.mu.Lock()
 	defer m.mu.Unlock()
@@ -562,13 +576,9 @@ type dataWriter struct {
 func (w *dataWriter) send(msg *fnpb.Elements) error {
 	if err := w.ch.client.Send(msg); err != nil {
 		if err == io.EOF {
-			log.Warnf(context.TODO(), "dataWriter[%v;%v] EOF on send; fetching real error", w.id, w.ch.id)
-			err = nil
-			for err == nil {
-				// Per GRPC stream documentation, if there's an EOF, we must call Recv
-				// until a non-nil error is returned, to ensure resources are cleaned up.
-				// https://pkg.go.dev/google.golang.org/grpc#ClientConn.NewStream
-				_, err = w.ch.client.Recv()
+			// Don't Recv here; the read loop owns the stream.
+			if w.ch.readErr != nil {
+				err = w.ch.readErr
 			}
 		}
 		log.Warnf(context.TODO(), "dataWriter[%v;%v] error on send: %v", w.id, w.ch.id, err)
@@ -687,13 +697,9 @@ type timerWriter struct {
 func (w *timerWriter) send(msg *fnpb.Elements) error {
 	if err := w.ch.client.Send(msg); err != nil {
 		if err == io.EOF {
-			log.Warnf(context.TODO(), "timerWriter[%v;%v] EOF on send; fetching real error", w.id, w.ch.id)
-			err = nil
-			for err == nil {
-				// Per GRPC stream documentation, if there's an EOF, we must call Recv
-				// until a non-nil error is returned, to ensure resources are cleaned up.
-				// https://pkg.go.dev/google.golang.org/grpc#ClientConn.NewStream
-				_, err = w.ch.client.Recv()
+			// Don't Recv here; the read loop owns the stream.
+			if w.ch.readErr != nil {
+				err = w.ch.readErr
 			}
 		}
 		log.Warnf(context.TODO(), "timerWriter[%v;%v] error on send: %v", w.id, w.ch.id, err)

sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go

@@ -24,6 +24,7 @@ import (
 	"runtime"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -543,6 +544,106 @@ func TestDataChannelTerminate_Writes(t *testing.T) {
 	}
 }
 
+type teardownDataClient struct{ block chan struct{} }
+
+func (f *teardownDataClient) Recv() (*fnpb.Elements, error) { <-f.block; return nil, io.EOF }
+func (f *teardownDataClient) Send(*fnpb.Elements) error     { return nil }
+
+func TestDataChannelManagerClose(t *testing.T) {
+	m := &DataChannelManager{ports: map[string]*DataChannel{}}
+	var cancelled [2]atomic.Bool
+	var unblock sync.Once
+	block := make(chan struct{})
+	for i := 0; i < 2; i++ {
+		i := i
+		ch := makeDataChannel(context.Background(), "port", &teardownDataClient{block: block}, func() {
+			cancelled[i].Store(true)
+			unblock.Do(func() { close(block) })
+		})
+		ch.forceRecreate = func(string, error) {
+			m.mu.Lock()
+			_ = m.ports
+			m.mu.Unlock()
+		}
+		m.ports["p"+string(rune('0'+i))] = ch
+	}
+
+	m.Close()
+
+	for i := range cancelled {
+		if !cancelled[i].Load() {
+			t.Errorf("channel %d not cancelled", i)
+		}
+	}
+	m.mu.Lock()
+	left := m.ports
+	m.mu.Unlock()
+	if left != nil {
+		t.Error("ports not cleared")
+	}
+	m.Close()
+}
+
+type eofOnSendClient struct {
+	recvForever chan struct{}
+}
+
+func (c *eofOnSendClient) Send(*fnpb.Elements) error { return io.EOF }
+func (c *eofOnSendClient) Recv() (*fnpb.Elements, error) {
+	<-c.recvForever
+	return nil, io.EOF
+}
+
+func TestDataWriterSendEOF(t *testing.T) {
+	ch := &DataChannel{
+		id:       "id",
+		client:   &eofOnSendClient{recvForever: make(chan struct{})},
+		cancelFn: func() {},
+	}
+	w := &dataWriter{ch: ch, id: clientID{ptransformID: "pt", instID: "inst"}}
+
+	done := make(chan error, 1)
+	go func() {
+		ch.mu.Lock()
+		defer ch.mu.Unlock()
+		done <- w.send(&fnpb.Elements{})
+	}()
+
+	select {
+	case err := <-done:
+		if err == nil {
+			t.Fatal("send succeeded")
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("send blocked")
+	}
+}
+
+func TestTimerWriterSendEOF(t *testing.T) {
+	ch := &DataChannel{
+		id:       "id",
+		client:   &eofOnSendClient{recvForever: make(chan struct{})},
+		cancelFn: func() {},
+	}
+	w := &timerWriter{ch: ch, id: clientID{ptransformID: "pt", instID: "inst"}, timerFamilyID: "fam"}
+
+	done := make(chan error, 1)
+	go func() {
+		ch.mu.Lock()
+		defer ch.mu.Unlock()
+		done <- w.send(&fnpb.Elements{})
+	}()
+
+	select {
+	case err := <-done:
+		if err == nil {
+			t.Fatal("send succeeded")
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("send blocked")
+	}
+}
+
 type noopDataClient struct {
 }
 

sdks/go/pkg/beam/core/runtime/harness/harness.go

@@ -164,6 +164,8 @@ func MainWithOptions(ctx context.Context, loggingEndpoint, controlEndpoint strin
 		runnerCapabilities:   rcMap,
 		elmTimeout:           elmTimeout,
 	}
+	defer ctrl.data.Close()
+	defer ctrl.state.Close()
 
 	if enabled, ok := rcMap[graphx.URNDataSampling]; ok && enabled {
 		ctrl.dataSampler = exec.NewDataSampler(ctx)

sdks/go/pkg/beam/core/runtime/harness/statemgr.go

@@ -627,6 +627,20 @@ func (m *StateChannelManager) Open(ctx context.Context, port exec.Port) (*StateC
 	return ch, nil
 }
 
+// Close closes all cached StateChannels.
+func (m *StateChannelManager) Close() {
+	m.mu.Lock()
+	chans := m.ports
+	m.ports = nil
+	m.mu.Unlock()
+	for _, ch := range chans {
+		ch.mu.Lock()
+		ch.forceRecreate = nil
+		ch.mu.Unlock()
+		ch.cancelFn()
+	}
+}
+
 type stateClient interface {
 	Send(*fnpb.StateRequest) error
 	Recv() (*fnpb.StateResponse, error)
@@ -704,7 +718,7 @@ func (c *StateChannel) read(ctx context.Context) {
 		msg, err := c.client.Recv()
 		if err != nil {
 			c.terminateStreamOnError(err)
-			if err == io.EOF {
+			if err == io.EOF || status.Code(err) == codes.Canceled {
 				log.Warnf(ctx, "StateChannel[%v].read: closed", c.id)
 				return
 			}
@@ -750,14 +764,13 @@ func (c *StateChannel) write(ctx context.Context) {
 	}
 
 	if err == io.EOF {
-		log.Warnf(ctx, "StateChannel[%v].write EOF on send; fetching real error", c.id)
-		err = nil
-		for err == nil {
-			// Per GRPC stream documentation, if there's an EOF, we must call Recv
-			// until a non-nil error is returned, to ensure resources are cleaned up.
-			// https://pkg.go.dev/google.golang.org/grpc#ClientConn.NewStream
-			_, err = c.client.Recv()
+		// Don't Recv here; the read loop owns the stream.
+		log.Warnf(ctx, "StateChannel[%v].write EOF on send", c.id)
+		c.mu.Lock()
+		if c.closedErr != nil {
+			err = c.closedErr
 		}
+		c.mu.Unlock()
 	}
 	log.Errorf(ctx, "StateChannel[%v].write error on send: %v", c.id, err)
 

sdks/go/pkg/beam/core/runtime/harness/statemgr_test.go

@@ -23,11 +23,14 @@ import (
 	"log"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/apache/beam/sdks/v2/go/pkg/beam/internal/errors"
 	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 // fakeStateClient replicates the call and response protocol
@@ -137,6 +140,21 @@ func TestStateChannel(t *testing.T) {
 			},
 			expectedErr:       io.EOF,
 			validateCancelled: true,
+		}, {
+			name: "readCanceled",
+			caseFn: func(t *testing.T, c *StateChannel, client *fakeStateClient) error {
+				go func() {
+					req := <-client.send
+					client.setRecvErr(status.Error(codes.Canceled, "context canceled"))
+					client.recv <- &fnpb.StateResponse{
+						Id: req.Id,
+					}
+				}()
+				_, err := c.Send(&fnpb.StateRequest{})
+				return err
+			},
+			expectedErr:       status.Error(codes.Canceled, "context canceled"),
+			validateCancelled: true,
 		}, {
 			name: "readOtherErr",
 			caseFn: func(t *testing.T, c *StateChannel, client *fakeStateClient) error {
@@ -179,18 +197,12 @@ func TestStateChannel(t *testing.T) {
 			caseFn: func(t *testing.T, c *StateChannel, client *fakeStateClient) error {
 				go func() {
 					client.setSendErr(io.EOF)
-					req := <-client.send
-					// This can be plumbed through on either side, write or read,
-					// the important part is that we get it.
-					client.setRecvErr(expectedError)
-					client.recv <- &fnpb.StateResponse{
-						Id: req.Id,
-					}
+					<-client.send
 				}()
 				_, err := c.Send(&fnpb.StateRequest{})
 				return err
 			},
-			expectedErr:       expectedError,
+			expectedErr:       io.EOF,
 			validateCancelled: true,
 		}, {
 			name: "writeOtherError",
@@ -500,6 +512,79 @@ func TestStateKeyWriter(t *testing.T) {
 	}
 }
 
+type teardownStateClient struct{ block chan struct{} }
+
+func (f *teardownStateClient) Recv() (*fnpb.StateResponse, error) {
+	<-f.block
+	return nil, status.Error(codes.Canceled, "context canceled")
+}
+func (f *teardownStateClient) Send(*fnpb.StateRequest) error { return nil }
+
+func TestStateChannelManagerClose(t *testing.T) {
+	m := &StateChannelManager{}
+	block := make(chan struct{})
+	var cancelled atomic.Bool
+	var unblock sync.Once
+	ch := makeStateChannel(context.Background(), "port", &teardownStateClient{block: block}, func() {
+		cancelled.Store(true)
+		unblock.Do(func() { close(block) })
+	})
+	ch.forceRecreate = func(string, error) {
+		m.mu.Lock()
+		_ = m.ports
+		m.mu.Unlock()
+	}
+	m.ports = map[string]*StateChannel{"p": ch}
+
+	m.Close()
+
+	if !cancelled.Load() {
+		t.Error("channel not cancelled")
+	}
+	m.mu.Lock()
+	left := m.ports
+	m.mu.Unlock()
+	if left != nil {
+		t.Error("ports not cleared")
+	}
+	m.Close()
+}
+
+type eofOnSendStateClient struct {
+	recvForever chan struct{}
+}
+
+func (c *eofOnSendStateClient) Send(*fnpb.StateRequest) error { return io.EOF }
+func (c *eofOnSendStateClient) Recv() (*fnpb.StateResponse, error) {
+	<-c.recvForever
+	return nil, io.EOF
+}
+
+func TestStateChannelWriteEOF(t *testing.T) {
+	c := &StateChannel{
+		id:        "id",
+		client:    &eofOnSendStateClient{recvForever: make(chan struct{})},
+		requests:  make(chan *fnpb.StateRequest, 1),
+		responses: make(map[string]chan<- *fnpb.StateResponse),
+		cancelFn:  func() {},
+		DoneCh:    make(chan struct{}),
+	}
+	c.responses["r1"] = make(chan *fnpb.StateResponse, 1)
+	c.requests <- &fnpb.StateRequest{Id: "r1"}
+
+	done := make(chan struct{})
+	go func() {
+		c.write(context.Background())
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("write blocked")
+	}
+}
+
 // This likely can't be replaced by the "errors" package helpers,
 // since we serialize errors in some cases.
 func contains(got, want error) bool {