Preserve NextRetryDelay across unrelated activity option updates

chasm/lib/activity/activity.go

@@ -23,6 +23,7 @@ import (
 	"fmt"
 	"math/rand"
 	"slices"
+	"strings"
 	"time"
 
 	"github.com/nexus-rpc/sdk-go/nexus"
@@ -639,13 +640,14 @@ func (a *Activity) UpdateActivityExecutionOptions(
 	}
 
 	frontendReq := req.GetFrontendRequest()
-
-	// start_delay updates are only valid while the activity is still in its delay window.
-	var hasStartDelayInMask bool
+	updateFields := map[string]struct{}{}
 	if mask := frontendReq.GetUpdateMask(); mask != nil {
-		_, hasStartDelayInMask = util.ParseFieldMask(mask)["startDelay"]
+		updateFields = util.ParseFieldMask(mask)
 	}
-	if !frontendReq.GetRestoreOriginal() && hasStartDelayInMask {
+
+	// start_delay updates are only valid while the activity is still in its delay window.
+	_, hasStartDelayInMask := updateFields["startDelay"]
+	if hasStartDelayInMask {
 		newDelay := frontendReq.GetActivityOptions().GetStartDelay()
 		if err := validateStartDelay(newDelay); err != nil {
 			return nil, err
@@ -663,6 +665,9 @@ func (a *Activity) UpdateActivityExecutionOptions(
 		}
 	}
 
+	attempt := a.LastAttempt.Get(ctx)
+	policyRetryIntervalBeforeUpdate := backoff.CalculateExponentialRetryInterval(a.RetryPolicy, attempt.GetCount()-1)
+
 	if frontendReq.GetRestoreOriginal() {
 		ogOptions := a.GetOriginalOptions()
 		a.TaskQueue = common.CloneProto(ogOptions.GetTaskQueue())
@@ -683,15 +688,14 @@ func (a *Activity) UpdateActivityExecutionOptions(
 		}
 	}
 
-	attempt := a.LastAttempt.Get(ctx)
-
-	// Recalculate the retry interval based on the (possibly updated) retry policy.
-	// This ensures a shortened retry interval takes effect immediately on re-dispatch.
-	status := a.GetStatus()
-	if (status == activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED ||
-		status == activitypb.ACTIVITY_EXECUTION_STATUS_PAUSED) &&
-		attempt.GetCurrentRetryInterval() != nil {
-
+	// Recalculate policy-derived retry intervals based on the (possibly updated) retry policy.
+	// Worker-provided NextRetryDelay values are preserved for their already-scheduled retry.
+	if a.shouldRecalculateCurrentRetryInterval(
+		attempt,
+		frontendReq.GetRestoreOriginal(),
+		updateFields,
+		policyRetryIntervalBeforeUpdate,
+	) {
 		newInterval := backoff.CalculateExponentialRetryInterval(a.RetryPolicy, attempt.GetCount()-1)
 		attempt.CurrentRetryInterval = durationpb.New(newInterval)
 	}
@@ -728,6 +732,52 @@ func (a *Activity) UpdateActivityExecutionOptions(
 	}, nil
 }
 
+// shouldRecalculateCurrentRetryInterval reports whether the pending retry's CurrentRetryInterval
+// should be re-derived from the (possibly updated) retry policy. All of the following must hold:
+//
+//   - the activity is waiting to retry, i.e. in SCHEDULED or PAUSED state (a running attempt has no
+//     pending backoff to recalculate);
+//   - a retry is actually pending (CurrentRetryInterval is set);
+//   - the update touches the retry policy: either RestoreOriginal (which replaces the whole policy),
+//     or the update mask includes "retryPolicy" or one of its subfields;
+//   - the pending interval is still policy-derived, i.e. it equals the pre-update policy value.
+//     A mismatch means the interval is a worker-provided NextRetryDelay override, which we preserve.
+func (a *Activity) shouldRecalculateCurrentRetryInterval(
+	attempt *activitypb.ActivityAttemptState,
+	restoreOriginal bool,
+	updateFields map[string]struct{},
+	policyRetryIntervalBeforeUpdate time.Duration,
+) bool {
+	status := a.GetStatus()
+	if status != activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED &&
+		status != activitypb.ACTIVITY_EXECUTION_STATUS_PAUSED {
+		return false
+	}
+
+	currentRetryInterval := attempt.GetCurrentRetryInterval()
+	if currentRetryInterval == nil {
+		return false
+	}
+
+	if !restoreOriginal {
+		hasRetryPolicyInMask := false
+		for field := range updateFields {
+			if field == "retryPolicy" || strings.HasPrefix(field, "retryPolicy.") {
+				hasRetryPolicyInMask = true
+				break
+			}
+		}
+		if !hasRetryPolicyInMask {
+			return false
+		}
+	}
+
+	// CurrentRetryInterval stores either policy-derived backoff or worker-provided NextRetryDelay overrides.
+	// Only recalculate intervals that match the old policy value; different values are treated as explicit
+	// worker overrides.
+	return currentRetryInterval.AsDuration() == policyRetryIntervalBeforeUpdate
+}
+
 // mergeActivityOptions applies the field mask from the request to the activity state.
 // The structure mirrors the field-mask logic in service/history/api/updateactivityoptions/api.go
 func (a *Activity) mergeActivityOptions(
@@ -897,6 +947,8 @@ func (a *Activity) handleUnpauseRequested(ctx chasm.MutableContext, req *activit
 		if err := TransitionUnpausedWhilePauseRequested.Apply(a, ctx, event); err != nil {
 			return nil, err
 		}
+	case activitypb.ACTIVITY_EXECUTION_STATUS_RESET_REQUESTED:
+		a.ResetKeepPaused = false
 	default:
 		return nil, serviceerror.NewFailedPreconditionf("activity is in non-unpausable state %v", a.GetStatus())
 	}
@@ -911,6 +963,8 @@ func (a *Activity) isPaused() bool {
 	case activitypb.ACTIVITY_EXECUTION_STATUS_PAUSED,
 		activitypb.ACTIVITY_EXECUTION_STATUS_PAUSE_REQUESTED:
 		return true
+	case activitypb.ACTIVITY_EXECUTION_STATUS_RESET_REQUESTED:
+		return a.GetResetKeepPaused()
 	default:
 		return false
 	}
@@ -921,17 +975,16 @@ func (a *Activity) unpause(
 	event unpauseEvent,
 ) {
 	attempt := a.LastAttempt.Get(ctx)
-	// Compute the dispatch time while the pending retry state is still intact; it is cleared below.
 	dispatchTime := a.unpauseDispatchTime(ctx, event)
 
 	if event.req.GetResetAttempts() {
 		attempt.Count = 1
+		attempt.CurrentRetryInterval = nil
 	}
 	if event.req.GetResetHeartbeat() {
 		a.LastHeartbeat = chasm.NewDataField(ctx, &activitypb.ActivityHeartbeatState{})
 	}
 	attempt.Stamp++
-	attempt.CurrentRetryInterval = nil
 	if timeout := a.GetScheduleToStartTimeout().AsDuration(); timeout > 0 {
 		ctx.AddTask(
 			a,

chasm/lib/activity/activity_tasks.go

@@ -137,8 +137,7 @@ func (h *scheduleToCloseTimeoutTaskHandler) Validate(
 	}
 	// Stamp check: discard tasks from before the most recent ScheduleToCloseTimeoutTask was
 	// scheduled (e.g. after a schedule-to-close extension or a disable+re-enable cycle).
-	// Tasks without a stamp (stamp=0) predate this field and are not validated by stamp.
-	if task.GetStamp() != 0 && task.GetStamp() != activity.GetScheduleToCloseStamp() {
+	if task.GetStamp() != activity.GetScheduleToCloseStamp() {
 		return false, nil
 	}
 	return true, nil

chasm/lib/activity/activity_test.go

@@ -643,3 +643,83 @@ func TestEffectiveUserMetadata_FallsBackToLegacy(t *testing.T) {
 	got := activity.effectiveUserMetadata(ctx)
 	require.Same(t, legacyMD, got)
 }
+
+func TestShouldRecalculateCurrentRetryInterval(t *testing.T) {
+	policyRetryInterval := 2 * time.Second
+
+	testCases := []struct {
+		name                 string
+		status               activitypb.ActivityExecutionStatus
+		restoreOriginal      bool
+		updateFields         map[string]struct{}
+		currentRetryInterval *durationpb.Duration
+		expectRecalculate    bool
+	}{
+		{
+			name:                 "retry policy subfield update with policy-derived interval",
+			status:               activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			updateFields:         map[string]struct{}{"retryPolicy.initialInterval": {}},
+			currentRetryInterval: durationpb.New(policyRetryInterval),
+			expectRecalculate:    true,
+		},
+		{
+			name:                 "retry policy replacement with policy-derived interval",
+			status:               activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			updateFields:         map[string]struct{}{"retryPolicy": {}},
+			currentRetryInterval: durationpb.New(policyRetryInterval),
+			expectRecalculate:    true,
+		},
+		{
+			name:                 "restore original with policy-derived interval",
+			status:               activitypb.ACTIVITY_EXECUTION_STATUS_PAUSED,
+			restoreOriginal:      true,
+			currentRetryInterval: durationpb.New(policyRetryInterval),
+			expectRecalculate:    true,
+		},
+		{
+			name:                 "unrelated update preserves retry interval",
+			status:               activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			updateFields:         map[string]struct{}{"heartbeatTimeout": {}},
+			currentRetryInterval: durationpb.New(policyRetryInterval),
+		},
+		{
+			name:                 "worker override differs from policy interval",
+			status:               activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			updateFields:         map[string]struct{}{"retryPolicy.initialInterval": {}},
+			currentRetryInterval: durationpb.New(policyRetryInterval + time.Second),
+		},
+		{
+			name:                 "started activity is not in retry backoff",
+			status:               activitypb.ACTIVITY_EXECUTION_STATUS_STARTED,
+			updateFields:         map[string]struct{}{"retryPolicy.initialInterval": {}},
+			currentRetryInterval: durationpb.New(policyRetryInterval),
+		},
+		{
+			name:         "missing retry interval",
+			status:       activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			updateFields: map[string]struct{}{"retryPolicy.initialInterval": {}},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			activity := &Activity{
+				ActivityState: &activitypb.ActivityState{
+					Status: tc.status,
+				},
+			}
+			attempt := &activitypb.ActivityAttemptState{
+				CurrentRetryInterval: tc.currentRetryInterval,
+			}
+
+			got := activity.shouldRecalculateCurrentRetryInterval(
+				attempt,
+				tc.restoreOriginal,
+				tc.updateFields,
+				policyRetryInterval,
+			)
+
+			require.Equal(t, tc.expectRecalculate, got)
+		})
+	}
+}

chasm/lib/activity/proto/v1/tasks.proto

@@ -17,7 +17,7 @@ message ScheduleToStartTimeoutTask {
 message ScheduleToCloseTimeoutTask {
   // The schedule-to-close stamp for this task. Used for task validation.
   // See also [ActivityState.schedule_to_close_stamp].
-  // Tasks without a stamp (stamp=0) predate this field and are not validated by stamp.
+  // Tasks without a stamp (stamp=0) predate this field and are valid only while the activity stamp is also 0.
   int32 stamp = 1;
 }