DON Time: clean ObservedDonTimes (#1904)

* Prune dontime store properly

* Update plugin.go

* Log state size

* Add tests

---------

Co-authored-by: Prashant Yadav <prashant.yadav@smartcontract.com>

Author: DylanTinianov

pkg/workflows/dontime/plugin.go

@@ -140,11 +140,11 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 		}
 
 		for id, requestSeqNum := range observation.Requests {
-			if _, ok := prevOutcome.ObservedDonTimes[id]; !ok {
-				prevOutcome.ObservedDonTimes[id] = &pb.ObservedDonTimes{}
+			var currSeqNum int64
+			if times, ok := prevOutcome.ObservedDonTimes[id]; ok {
+				currSeqNum = int64(len(times.Timestamps))
 			}
 			// We only count requests for the next sequence number and ignore all other ones.
-			currSeqNum := int64(len(prevOutcome.ObservedDonTimes[id].Timestamps))
 			if requestSeqNum == currSeqNum {
 				observationCounts[id]++
 			} else if requestSeqNum > currSeqNum {
@@ -197,17 +197,25 @@ func (p *Plugin) Outcome(_ context.Context, outctx ocr3types.OutcomeContext, _ t
 		}
 	}
 
-	// Remove expired workflow executions
+	// Remove expired and empty workflow executions
 	for id, observedTimes := range outcome.ObservedDonTimes {
-		if observedTimes != nil && len(observedTimes.Timestamps) > 0 {
-			if donTime >= observedTimes.Timestamps[0]+p.offChainConfig.ExecutionRemovalTime.AsDuration().Milliseconds() {
-				delete(outcome.ObservedDonTimes, id)
-				p.store.deleteExecutionID(id)
-			}
+		if observedTimes == nil || len(observedTimes.Timestamps) == 0 {
+			delete(outcome.ObservedDonTimes, id)
+			p.store.deleteExecutionID(id)
+			continue
+		}
+		if donTime >= observedTimes.Timestamps[0]+p.offChainConfig.ExecutionRemovalTime.AsDuration().Milliseconds() {
+			delete(outcome.ObservedDonTimes, id)
+			p.store.deleteExecutionID(id)
 		}
 	}
 
-	return proto.MarshalOptions{Deterministic: true}.Marshal(outcome)
+	outcomeBytes, err := proto.MarshalOptions{Deterministic: true}.Marshal(outcome)
+	p.lggr.Infow("Outcome computed",
+		"observedDonTimesEntries", len(outcome.ObservedDonTimes),
+		"outcomeSizeBytes", len(outcomeBytes),
+	)
+	return outcomeBytes, err
 }
 
 func (p *Plugin) Reports(_ context.Context, _ uint64, outcome ocr3types.Outcome) ([]ocr3types.ReportPlus[[]byte], error) {

pkg/workflows/dontime/plugin_test.go

@@ -213,6 +213,210 @@ func TestPlugin_Outcome(t *testing.T) {
 	require.Equal(t, []int64{timestamp}, outcomeProto.ObservedDonTimes[executionID].Timestamps)
 }
 
+func TestPlugin_Outcome_SequenceNumberHandling(t *testing.T) {
+	lggr := logger.Test(t)
+	config, offchainCfg := newTestPluginConfig(t), newTestPluginOffchainConfig(t)
+	ctx := t.Context()
+
+	makeObservations := func(t *testing.T, timestamp int64, requests map[string]int64, numNodes int) []types.AttributedObservation {
+		t.Helper()
+		aos := make([]types.AttributedObservation, numNodes)
+		for i := 0; i < numNodes; i++ {
+			obs := &pb.Observation{
+				Timestamp: timestamp + int64(i),
+				Requests:  requests,
+			}
+			rawObs, err := proto.Marshal(obs)
+			require.NoError(t, err)
+			aos[i] = types.AttributedObservation{
+				Observation: rawObs,
+				Observer:    commontypes.OracleID(i),
+			}
+		}
+		return aos
+	}
+
+	t.Run("new execution ID not in previous outcome defaults currSeqNum to 0", func(t *testing.T) {
+		store := NewStore(DefaultRequestTimeout)
+		plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+		require.NoError(t, err)
+
+		executionID := "new-workflow"
+		_ = store.RequestDonTime(executionID, 0)
+
+		timestamp := time.Now().UnixMilli()
+		aos := makeObservations(t, timestamp, map[string]int64{executionID: 0}, 4)
+
+		prevOutcome := &pb.Outcome{
+			Timestamp:        0,
+			ObservedDonTimes: map[string]*pb.ObservedDonTimes{},
+		}
+		prevOutcomeBytes, err := proto.Marshal(prevOutcome)
+		require.NoError(t, err)
+
+		query, err := plugin.Query(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes})
+		require.NoError(t, err)
+
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, aos)
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+
+		require.Contains(t, outcomeProto.ObservedDonTimes, executionID)
+		require.Len(t, outcomeProto.ObservedDonTimes[executionID].Timestamps, 1)
+	})
+
+	t.Run("nil ObservedDonTimes in previous outcome does not panic", func(t *testing.T) {
+		store := NewStore(DefaultRequestTimeout)
+		plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+		require.NoError(t, err)
+
+		executionID := "nil-map-workflow"
+		_ = store.RequestDonTime(executionID, 0)
+
+		timestamp := time.Now().UnixMilli()
+		aos := makeObservations(t, timestamp, map[string]int64{executionID: 0}, 4)
+
+		prevOutcome := &pb.Outcome{
+			Timestamp:        0,
+			ObservedDonTimes: nil,
+		}
+		prevOutcomeBytes, err := proto.Marshal(prevOutcome)
+		require.NoError(t, err)
+
+		query, err := plugin.Query(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes})
+		require.NoError(t, err)
+
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, aos)
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+
+		require.Contains(t, outcomeProto.ObservedDonTimes, executionID)
+		require.Len(t, outcomeProto.ObservedDonTimes[executionID].Timestamps, 1)
+	})
+
+	t.Run("existing execution ID uses len(Timestamps) as currSeqNum", func(t *testing.T) {
+		store := NewStore(DefaultRequestTimeout)
+		plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+		require.NoError(t, err)
+
+		executionID := "existing-workflow"
+		_ = store.RequestDonTime(executionID, 1)
+
+		timestamp := time.Now().UnixMilli()
+		aos := makeObservations(t, timestamp, map[string]int64{executionID: 1}, 4)
+
+		prevTimestamp := timestamp - 1000 // 1 second ago in millis
+		prevOutcome := &pb.Outcome{
+			Timestamp: prevTimestamp,
+			ObservedDonTimes: map[string]*pb.ObservedDonTimes{
+				executionID: {Timestamps: []int64{prevTimestamp}},
+			},
+		}
+		prevOutcomeBytes, err := proto.Marshal(prevOutcome)
+		require.NoError(t, err)
+
+		query, err := plugin.Query(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes})
+		require.NoError(t, err)
+
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, aos)
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+
+		require.Contains(t, outcomeProto.ObservedDonTimes, executionID)
+		require.Len(t, outcomeProto.ObservedDonTimes[executionID].Timestamps, 2)
+	})
+
+	t.Run("stale sequence number is ignored", func(t *testing.T) {
+		store := NewStore(DefaultRequestTimeout)
+		plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+		require.NoError(t, err)
+
+		executionID := "stale-workflow"
+
+		timestamp := time.Now().UnixMilli()
+		// Observations report seqNum 0, but prevOutcome already has 2 timestamps (currSeqNum=2)
+		aos := makeObservations(t, timestamp, map[string]int64{executionID: 0}, 4)
+
+		prevTimestamp := timestamp - 1000 // 1 second ago in millis
+		prevOutcome := &pb.Outcome{
+			Timestamp: prevTimestamp,
+			ObservedDonTimes: map[string]*pb.ObservedDonTimes{
+				executionID: {Timestamps: []int64{
+					prevTimestamp - 1000,
+					prevTimestamp,
+				}},
+			},
+		}
+		prevOutcomeBytes, err := proto.Marshal(prevOutcome)
+		require.NoError(t, err)
+
+		query, err := plugin.Query(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes})
+		require.NoError(t, err)
+
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, aos)
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+
+		// Stale seqNum 0 should be ignored, so timestamps should remain unchanged at 2
+		require.Len(t, outcomeProto.ObservedDonTimes[executionID].Timestamps, 2)
+	})
+
+	t.Run("mix of new and existing execution IDs", func(t *testing.T) {
+		store := NewStore(DefaultRequestTimeout)
+		plugin, err := NewPlugin(store, config, offchainCfg, lggr)
+		require.NoError(t, err)
+
+		existingID := "existing-workflow"
+		newID := "new-workflow"
+		_ = store.RequestDonTime(existingID, 1)
+		_ = store.RequestDonTime(newID, 0)
+
+		timestamp := time.Now().UnixMilli()
+		requests := map[string]int64{
+			existingID: 1,
+			newID:      0,
+		}
+		aos := makeObservations(t, timestamp, requests, 4)
+
+		prevTimestamp := timestamp - 1000 // 1 second ago in millis
+		prevOutcome := &pb.Outcome{
+			Timestamp: prevTimestamp,
+			ObservedDonTimes: map[string]*pb.ObservedDonTimes{
+				existingID: {Timestamps: []int64{prevTimestamp}},
+			},
+		}
+		prevOutcomeBytes, err := proto.Marshal(prevOutcome)
+		require.NoError(t, err)
+
+		query, err := plugin.Query(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes})
+		require.NoError(t, err)
+
+		outcome, err := plugin.Outcome(ctx, ocr3types.OutcomeContext{PreviousOutcome: prevOutcomeBytes}, query, aos)
+		require.NoError(t, err)
+
+		outcomeProto := &pb.Outcome{}
+		err = proto.Unmarshal(outcome, outcomeProto)
+		require.NoError(t, err)
+
+		require.Contains(t, outcomeProto.ObservedDonTimes, existingID)
+		require.Len(t, outcomeProto.ObservedDonTimes[existingID].Timestamps, 2)
+		require.Contains(t, outcomeProto.ObservedDonTimes, newID)
+		require.Len(t, outcomeProto.ObservedDonTimes[newID].Timestamps, 1)
+	})
+}
+
 func TestPlugin_FinishedExecutions(t *testing.T) {
 	lggr := logger.Test(t)
 	store := NewStore(DefaultRequestTimeout)
@@ -280,11 +484,16 @@ func TestPlugin_FinishedExecutions(t *testing.T) {
 	})
 
 	t.Run("Transmit: delete removed executionIDs", func(t *testing.T) {
+		store.setDonTimes("workflow-123", []int64{time.Now().UnixMilli()})
+
 		r := ocr3types.ReportWithInfo[[]byte]{}
 		r.Report, err = proto.Marshal(outcomeProto)
 		require.NoError(t, err)
 		err = transmitter.Transmit(ctx, types.ConfigDigest{}, 0, r, []types.AttributedOnchainSignature{})
 		require.NoError(t, err)
+
+		_, err = store.GetDonTimes("workflow-123")
+		require.ErrorContains(t, err, "no don time for executionID workflow-123")
 	})
 }
 

pkg/workflows/dontime/store.go

@@ -117,6 +117,22 @@ func (s *Store) setDonTimes(executionID string, donTimes []int64) {
 	s.donTimes[executionID] = donTimes
 }
 
+func (s *Store) replaceDonTimes(donTimes map[string][]int64) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	for executionID, timestamps := range donTimes {
+		s.donTimes[executionID] = timestamps
+	}
+
+	for executionID := range s.donTimes {
+		if _, ok := donTimes[executionID]; !ok {
+			delete(s.donTimes, executionID)
+			delete(s.requests, executionID)
+		}
+	}
+}
+
 func (s *Store) GetLastObservedDonTime() int64 {
 	s.mu.Lock()
 	defer s.mu.Unlock()

pkg/workflows/dontime/transmitter.go

@@ -33,9 +33,11 @@ func (t *Transmitter) Transmit(_ context.Context, _ types.ConfigDigest, _ uint64
 		return err
 	}
 
+	currentDonTimes := make(map[string][]int64, len(outcome.ObservedDonTimes))
 	for id, observedDonTimes := range outcome.ObservedDonTimes {
-		t.store.setDonTimes(id, observedDonTimes.Timestamps)
+		currentDonTimes[id] = observedDonTimes.Timestamps
 	}
+	t.store.replaceDonTimes(currentDonTimes)
 	t.store.setLastObservedDonTime(outcome.Timestamp)
 
 	t.lggr.Infow("Transmitting timestamps", "lastObservedDonTime", outcome.Timestamp)