This is an automated email from the ASF dual-hosted git repository.
lostluck pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new de688c0a199 Close harness data/state channels and don't Recv on Send
EOF (#40262)
de688c0a199 is described below
commit de688c0a19961f3bbf86f8ae594e7a243c1a8d59
Author: Igor <[email protected]>
AuthorDate: Thu Oct 1 13:10:19 2026 +0300
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 <[email protected]>
---
CHANGES.md | 1 +
sdks/go/pkg/beam/core/runtime/harness/datamgr.go | 34 ++++---
.../pkg/beam/core/runtime/harness/datamgr_test.go | 101 +++++++++++++++++++++
sdks/go/pkg/beam/core/runtime/harness/harness.go | 2 +
sdks/go/pkg/beam/core/runtime/harness/statemgr.go | 29 ++++--
.../pkg/beam/core/runtime/harness/statemgr_test.go | 101 +++++++++++++++++++--
6 files changed, 238 insertions(+), 30 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 5ef4f34c92b..67db578fc27 100644
--- a/CHANGES.md
+++ b/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)).
diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
index ab93dcba51b..b9718e9b42b 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
+++ b/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)
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 9f6f8a986a3..637076d6e44 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go
+++ b/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 {
}
diff --git a/sdks/go/pkg/beam/core/runtime/harness/harness.go
b/sdks/go/pkg/beam/core/runtime/harness/harness.go
index 969ac1b0a64..45a9f05d736 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/harness.go
+++ b/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)
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statemgr.go
b/sdks/go/pkg/beam/core/runtime/harness/statemgr.go
index 269ded37299..09695ac2ef4 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/statemgr.go
+++ b/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)
diff --git a/sdks/go/pkg/beam/core/runtime/harness/statemgr_test.go
b/sdks/go/pkg/beam/core/runtime/harness/statemgr_test.go
index 53c4ca05ee1..5416cdfad72 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/statemgr_test.go
+++ b/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 {
