[Prism] Fix potential side-effect in TriggerAfterEach.onFire. (#36166)

* Fix potential side-effect in TriggerAfterEach.onFire.

* Add a unit test.

Author: shunping

sdks/go/pkg/beam/runners/prism/internal/engine/strategy.go

@@ -302,15 +302,23 @@ func (t *TriggerAfterEach) onFire(state *StateData) {
 	if !t.shouldFire(state) {
 		return
 	}
-	for _, sub := range t.SubTriggers {
+	for i, sub := range t.SubTriggers {
 		if state.getTriggerState(sub).finished {
 			continue
 		}
 		sub.onFire(state)
+		// If the sub-trigger didn't finish, we return, waiting for it to finish on a subsequent call.
 		if !state.getTriggerState(sub).finished {
 			return
 		}
+
+		// If the sub-trigger finished, we check if it's the last one.
+		// If it's not the last one, we return, waiting for the next onFire call to advance to the next sub-trigger.
+		if i < len(t.SubTriggers)-1 {
+			return
+		}
 	}
+	// clear and reset when all sub-triggers have fired.
 	triggerClearAndFinish(t, state)
 }
 

sdks/go/pkg/beam/runners/prism/internal/engine/strategy_test.go

@@ -122,6 +122,25 @@ func TestTriggers_isReady(t *testing.T) {
 				{triggerInput{newElementCount: 1}, false},
 				{triggerInput{newElementCount: 1}, false},
 			},
+		}, {
+			name: "afterEach_2_Always_1",
+			trig: &TriggerAfterEach{
+				SubTriggers: []Trigger{
+					&TriggerElementCount{2},
+					&TriggerAfterAny{SubTriggers: []Trigger{&TriggerAlways{}}},
+					&TriggerElementCount{1},
+				},
+			},
+			inputs: []io{
+				{triggerInput{newElementCount: 1}, false},
+				{triggerInput{newElementCount: 1}, true},  // first is ready
+				{triggerInput{newElementCount: 1}, true},  // second is ready
+				{triggerInput{newElementCount: 1}, true},  // third is ready
+				{triggerInput{newElementCount: 1}, false}, // never resets after this.
+				{triggerInput{newElementCount: 1}, false},
+				{triggerInput{newElementCount: 1}, false},
+				{triggerInput{newElementCount: 1}, false},
+			},
 		}, {
 			name: "afterAny_2_3_4",
 			trig: &TriggerAfterAny{