(beam) branch master updated: [#29605] Go SDK: Eagerly create timer writers to ensure is_last sent, minmize lock contention. (#29607)

[lo...@apache.org]

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 e19ace02a28 [#29605] Go SDK: Eagerly create timer writers to ensure is_last sent, minmize lock contention. (#29607)
e19ace02a28 is described below

commit e19ace02a28f4fa19ca3db7f0dc43e9c4f4ab3c1
Author: Robert Burke <lo...@users.noreply.github.com>
AuthorDate: Fri Dec 8 10:36:25 2023 -0800

    [#29605] Go SDK: Eagerly create timer writers to ensure is_last sent, minmize lock contention. (#29607)
---
 sdks/go/examples/timer_wordcap/wordcap.go          |  1 -
 sdks/go/pkg/beam/core/runtime/harness/datamgr.go   | 23 ++++++++++++++++++++++
 .../pkg/beam/core/runtime/harness/datamgr_test.go  | 12 +++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/sdks/go/examples/timer_wordcap/wordcap.go b/sdks/go/examples/timer_wordcap/wordcap.go
index db64c10eb4d..01f87edc59b 100644
--- a/sdks/go/examples/timer_wordcap/wordcap.go
+++ b/sdks/go/examples/timer_wordcap/wordcap.go
@@ -118,7 +118,6 @@ func (s *Stateful) OnTimer(ctx context.Context, ts beam.EventTime, sp state.Prov
 			// Clean up the state that has been evicted.
 			s.ElementBag.Clear(sp)
 			s.MinTime.Clear(sp)
-			s.OutputState.ClearTag(tp, tag) // Clean up the fired timer tag. (Temporary workaround for a runner bug.)
 		}
 	}
 }
diff --git a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
index ed57e3eca59..0f2de99dd2a 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/datamgr.go
@@ -386,14 +386,31 @@ func (c *DataChannel) read(ctx context.Context) {
 			return
 		}
 
+		// Consolidating required timer writer creation to a optional single lock section.
+		type seenTimers struct {
+			InstID                 instructionID
+			PTransformID, FamilyID string
+		}
+		neededTimers := map[seenTimers]struct{}{}
+
 		// Each message may contain segments for multiple streams, so we
 		// must treat each segment in isolation. We maintain a local cache
 		// to reduce lock contention.
 		iterateElements(c, cache, &seenLast, msg.GetTimers(),
 			func(elm *fnpb.Elements_Timers) exec.Elements {
+				neededTimers[seenTimers{InstID: instructionID(elm.GetInstructionId()), PTransformID: elm.GetTransformId(), FamilyID: elm.GetTimerFamilyId()}] = struct{}{}
 				return exec.Elements{Timers: elm.GetTimers(), PtransformID: elm.GetTransformId(), TimerFamilyID: elm.GetTimerFamilyId()}
 			})
 
+		// Creating a writer is necessary to ensure a "is_last" signal is returned for timers that aren't set.
+		if len(neededTimers) > 0 {
+			c.mu.Lock()
+			for key := range neededTimers {
+				c.makeTimerWriterLocked(ctx, clientID{ptransformID: key.PTransformID, instID: key.InstID}, key.FamilyID)
+			}
+			c.mu.Unlock()
+		}
+
 		iterateElements(c, cache, &seenLast, msg.GetData(),
 			func(elm *fnpb.Elements_Data) exec.Elements {
 				return exec.Elements{Data: elm.GetData(), PtransformID: elm.GetTransformId()}
@@ -629,7 +646,13 @@ func (w *dataWriter) Write(p []byte) (n int, err error) {
 func (c *DataChannel) makeTimerWriter(ctx context.Context, id clientID, family string) *timerWriter {
 	c.mu.Lock()
 	defer c.mu.Unlock()
+	return c.makeTimerWriterLocked(ctx, id, family)
+}
 
+// makeTimerWriterLocked does the work of makeTimerWriter, but doesn't call the lock methods.
+//
+// c.mu must be locked when this is called.
+func (c *DataChannel) makeTimerWriterLocked(ctx context.Context, id clientID, family string) *timerWriter {
 	var m map[timerKey]*timerWriter
 	var ok bool
 	if m, ok = c.timerWriters[id.instID]; !ok {
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 c7f8ac5858c..92c4d0a8f8c 100644
--- a/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go
+++ b/sdks/go/pkg/beam/core/runtime/harness/datamgr_test.go
@@ -389,6 +389,18 @@ func TestElementChan(t *testing.T) {
 				return elms
 			},
 			wantSum: 0, wantCount: 0,
+		}, {
+			name: "SomeTimersAndADataThenReaderThenCleanup",
+			sequenceFn: func(ctx context.Context, t *testing.T, client *fakeChanClient, c *DataChannel) <-chan exec.Elements {
+				client.Send(&fnpb.Elements{
+					Timers: []*fnpb.Elements_Timers{timerElm(1, false), timerElm(2, true)},
+					Data:   []*fnpb.Elements_Data{dataElm(3, true)},
+				})
+				elms := openChan(ctx, t, c, timerID)
+				c.removeInstruction(instID)
+				return elms
+			},
+			wantSum: 6, wantCount: 3,
 		},
 	}
 	for _, test := range tests {