Fix deadlock on worker failures by introducing a bundle-level Done signal

The cancellation signal `b.Done` is to abort bundle data streaming immediately
if the worker fails or completes the bundle early.

Without this signal, a deadlock scenario can occur.
- Runner sends elements through data channel
- SDK worker reads and processes the first element, raises an exception
  and sends InstructionResponse through control channel. However,
  at the same time, Runner keeps sending data and is blocked until
  SDK worker reads more data.
- The existed two signals are not sufficient to do an immediate break
  - ctx.Done() is closed when the pipeline is done or timeout. Given
    the bundle failure, timeout is the only option but it is too late.
  - wk.StoppedChan() is closed when the worker is stopped by the worker
    pool, which is also not happening when the runner is waiting to
    send data.

Author: shunping

sdks/go/pkg/beam/runners/prism/internal/worker/bundle.go

@@ -62,6 +62,7 @@ type B struct {
 	OutputData engine.TentativeData
 
 	Resp      chan *fnpb.ProcessBundleResponse
+	Done      chan struct{}
 	BundleErr error
 	responded bool
 
@@ -80,6 +81,7 @@ func (b *B) Init() {
 		close(b.DataWait) // Can happen if there are no outputs for the bundle.
 	}
 	b.Resp = make(chan *fnpb.ProcessBundleResponse, 1)
+	b.Done = make(chan struct{})
 }
 
 // DataOrTimerDone indicates a final element has been received from a Data or Timer output.
@@ -102,7 +104,11 @@ func (b *B) Respond(resp *fnpb.InstructionResponse) {
 		return
 	}
 	b.responded = true
+	if b.Done != nil {
+		close(b.Done)
+	}
 	if resp.GetError() != "" {
+		slog.Error("DEBUG: Prism received bundle error from worker response", "bundle", resp.GetInstructionId())
 		b.BundleErr = fmt.Errorf("bundle %v %v failed:%v", resp.GetInstructionId(), b.PBDID, resp.GetError())
 		close(b.Resp)
 		return
@@ -143,6 +149,13 @@ func (b *B) ProcessOn(ctx context.Context, wk *W) <-chan struct{} {
 			b.DataOrTimerDone()
 		}
 		return b.DataWait
+	case <-b.Done:
+		// The bundle completed/failed before req was sent.
+		outCap := b.OutputCount + len(b.HasTimers)
+		for i := 0; i < outCap; i++ {
+			b.DataOrTimerDone()
+		}
+		return b.DataWait
 	case wk.InstReqs <- req:
 		// desired outcome
 	}
@@ -181,6 +194,9 @@ func (b *B) ProcessOn(ctx context.Context, wk *W) <-chan struct{} {
 		case <-ctx.Done():
 			b.DataOrTimerDone()
 			return b.DataWait
+		case <-b.Done:
+			b.DataOrTimerDone()
+			return b.DataWait
 		case wk.DataReqs <- elms:
 		}
 	}
@@ -202,6 +218,9 @@ func (b *B) ProcessOn(ctx context.Context, wk *W) <-chan struct{} {
 	case <-ctx.Done():
 		b.DataOrTimerDone()
 		return b.DataWait
+	case <-b.Done:
+		b.DataOrTimerDone()
+		return b.DataWait
 	case wk.DataReqs <- &fnpb.Elements{
 		Timers: timers,
 		Data: []*fnpb.Elements_Data{

sdks/go/pkg/beam/runners/prism/internal/worker/worker_test.go

@@ -522,6 +522,73 @@ func TestWorker_State_MultimapSideInput(t *testing.T) {
 	}
 }
 
+// TestBundle_ProcessOn_WorkerFailure verifies that the runner does not deadlock when
+// a worker fails mid-bundle and stops reading elements from the Data plane stream.
+func TestBundle_ProcessOn_WorkerFailure(t *testing.T) {
+	ctx, wk, clientConn := serveTestWorker(t)
+
+	dataCli := fnpb.NewBeamFnDataClient(clientConn)
+	dataStream, err := dataCli.Data(ctx)
+	if err != nil {
+		t.Fatal("couldn't create data client:", err)
+	}
+
+	instID := wk.NextInst()
+
+	// Create 15 large input blocks (512 KB each) to saturate the 10-slot channel buffer
+	// and the gRPC flow control window, forcing the Data sender inside worker.go to block.
+	largeBytes := make([]byte, 512*1024)
+	var inputBlocks []*engine.Block
+	for i := 0; i < 15; i++ {
+		inputBlocks = append(inputBlocks, &engine.Block{
+			Kind:  engine.BlockData,
+			Bytes: [][]byte{largeBytes},
+		})
+	}
+
+	b := &B{
+		InstID:      instID,
+		PBDID:       "teststageID",
+		Input:       inputBlocks,
+		OutputCount: 1,
+	}
+	b.Init()
+	wk.activeInstructions[instID] = b
+
+	processOnDone := make(chan struct{})
+	go func() {
+		b.ProcessOn(ctx, wk)
+		close(processOnDone)
+	}()
+
+	// Send the initial process bundle request trigger.
+	wk.InstReqs <- &fnpb.InstructionRequest{
+		InstructionId: instID,
+	}
+
+	// Read only the first block to simulate worker processing start.
+	_, err = dataStream.Recv()
+	if err != nil {
+		t.Fatal("couldn't receive first data element:", err)
+	}
+
+	// Simulate worker failure by responding with an error on the Control channel.
+	// Without the fix, ProcessOn's background goroutine deadlocks at `wk.DataReqs <- elms`
+	// because the client stopped reading and the buffer/flow-control is saturated.
+	wk.activeInstructions[instID].Respond(&fnpb.InstructionResponse{
+		InstructionId: instID,
+		Error:         "Intentional worker failure",
+	})
+
+	// Verify that ProcessOn exits cleanly and does not deadlock/hang.
+	select {
+	case <-processOnDone:
+		// Test passed: ProcessOn exited successfully!
+	case <-time.After(10 * time.Second):
+		t.Fatal("ProcessOn deadlocked / hung after worker failure!")
+	}
+}
+
 func newWorker() *W {
 	mw := &MultiplexW{
 		pool: map[string]*W{},