tasks/chasm/lib/scheduler/migration/migration.go

package migration

import (
	"time"

	commonpb "go.temporal.io/api/common/v1"
	enumspb "go.temporal.io/api/enums/v1"
	failurepb "go.temporal.io/api/failure/v1"
	schedulepb "go.temporal.io/api/schedule/v1"
	schedulespb "go.temporal.io/server/api/schedule/v1"
	schedulerpb "go.temporal.io/server/chasm/lib/scheduler/gen/schedulerpb/v1"
	"go.temporal.io/server/common"
	schedulescommon "go.temporal.io/server/common/schedules"
	"google.golang.org/protobuf/types/known/timestamppb"
)

// LegacyToCreateFromMigrationStateRequest converts legacy (workflow-backed) scheduler
// state to a CreateFromMigrationStateRequest proto. This is the primary V1-to-V2
// migration function.
//
// The migrationTime parameter is used for initializing timestamps that don't have a
// direct mapping from V1 state (e.g., StartTime for running workflows).
//
// State preserved during migration:
//   - Schedule spec, action, policies, and state
//   - ScheduleInfo metadata (action counts, create/update times)
//   - ConflictToken (for optimistic concurrency control)
//   - Buffered starts (with V2-specific fields like RequestId, WorkflowId populated)
//   - Running workflows (converted to BufferedStarts with RunId/StartTime set)
//   - Recent actions (converted to BufferedStarts with Completed field set)
//   - Ongoing backfills (converted to Backfiller components)
//   - Last completion result and continued failure
//   - High water mark (becomes Generator.LastProcessedTime)
//   - Search attributes and memo
//
// Note: In V2, RunningWorkflows and RecentActions are computed on-demand from
// BufferedStarts by the Invoker, rather than being stored separately in ScheduleInfo.
func LegacyToCreateFromMigrationStateRequest(
	schedule *schedulepb.Schedule,
	info *schedulepb.ScheduleInfo,
	state *schedulespb.InternalState,
	searchAttributes *commonpb.SearchAttributes,
	memo *commonpb.Memo,
	migrationTime time.Time,
) *schedulerpb.CreateFromMigrationStateRequest {
	// V2 computes RunningWorkflows/RecentActions on-demand from BufferedStarts
	infoClone := common.CloneProto(info)
	infoClone.RunningWorkflows = nil
	infoClone.RecentActions = nil

	schedulerState := &schedulerpb.SchedulerState{
		Schedule:      common.CloneProto(schedule),
		Info:          infoClone,
		Namespace:     state.Namespace,
		NamespaceId:   state.NamespaceId,
		ScheduleId:    state.ScheduleId,
		ConflictToken: state.ConflictToken,
		Closed:        false,
	}

	generatorState := &schedulerpb.GeneratorState{
		LastProcessedTime: common.CloneProto(state.LastProcessedTime),
		FutureActionTimes: nil, // Regenerated by GeneratorTask
	}

	pendingBufferedStarts := convertBufferedStartsLegacyToCHASM(
		state.BufferedStarts,
		state.NamespaceId,
		state.ScheduleId,
		state.ConflictToken,
		getWorkflowID(schedule),
	)

	runningBufferedStarts := convertRunningWorkflowsToBufferedStarts(
		info.RunningWorkflows,
		state.NamespaceId,
		state.ScheduleId,
		state.ConflictToken,
		migrationTime,
	)

	recentActionsBufferedStarts := convertRecentActionsToBufferedStarts(
		info.RecentActions,
		info.RunningWorkflows,
		state.NamespaceId,
		state.ScheduleId,
		state.ConflictToken,
		migrationTime,
	)

	allBufferedStarts := append(pendingBufferedStarts, runningBufferedStarts...)
	allBufferedStarts = append(allBufferedStarts, recentActionsBufferedStarts...)

	invokerState := &schedulerpb.InvokerState{
		BufferedStarts:    allBufferedStarts,
		LastProcessedTime: timestamppb.New(migrationTime),
	}

	backfillers := convertBackfillsLegacyToCHASM(state.OngoingBackfills)
	lastCompletion := convertLastCompletionLegacyToCHASM(state.LastCompletionResult, state.ContinuedFailure)

	return &schedulerpb.CreateFromMigrationStateRequest{
		NamespaceId: state.NamespaceId,
		State: &schedulerpb.SchedulerMigrationState{
			SchedulerState:       schedulerState,
			GeneratorState:       generatorState,
			InvokerState:         invokerState,
			Backfillers:          backfillers,
			LastCompletionResult: lastCompletion,
			SearchAttributes:     searchAttributes.GetIndexedFields(),
			Memo:                 memo.GetFields(),
		},
	}
}

// CHASMToLegacyStartScheduleArgs converts CHASM scheduler state to V1 StartScheduleArgs.
// This is the primary V2-to-V1 migration function. The migrationTime parameter is used
// to initialize missing timestamps.
func CHASMToLegacyStartScheduleArgs(
	scheduler *schedulerpb.SchedulerState,
	generator *schedulerpb.GeneratorState,
	invoker *schedulerpb.InvokerState,
	backfillers map[string]*schedulerpb.BackfillerState,
	lastCompletionResult *schedulerpb.LastCompletionResult,
	searchAttributes map[string]*commonpb.Payload,
	memo map[string]*commonpb.Payload,
	migrationTime time.Time,
) *schedulespb.StartScheduleArgs {
	schedulerState := common.CloneProto(scheduler)
	if schedulerState == nil {
		schedulerState = &schedulerpb.SchedulerState{}
	}

	schedule := common.CloneProto(schedulerState.Schedule)
	if schedule == nil {
		schedule = &schedulepb.Schedule{}
	}

	info := common.CloneProto(schedulerState.Info)
	if info == nil {
		info = &schedulepb.ScheduleInfo{}
	}

	var invokerBuffered []*schedulespb.BufferedStart
	if invoker != nil {
		invokerBuffered = invoker.GetBufferedStarts()
	}
	bufferedStarts, running, recent := splitBufferedStartsForLegacy(invokerBuffered)
	ongoingBackfills, triggerStarts := convertBackfillersCHASMToLegacy(backfillers, migrationTime)
	bufferedStarts = append(bufferedStarts, triggerStarts...)

	var generatorLastProcessed *timestamppb.Timestamp
	if generator != nil {
		generatorLastProcessed = generator.GetLastProcessedTime()
	}
	lastProcessedTime := common.CloneProto(generatorLastProcessed)
	if lastProcessedTime == nil {
		lastProcessedTime = timestamppb.New(migrationTime)
	}

	resultPayloads, continuedFailure := convertLastCompletionCHASMToLegacy(lastCompletionResult)

	info.RunningWorkflows = running
	info.RecentActions = recent

	state := &schedulespb.InternalState{
		Namespace:            schedulerState.Namespace,
		NamespaceId:          schedulerState.NamespaceId,
		ScheduleId:           schedulerState.ScheduleId,
		LastProcessedTime:    lastProcessedTime,
		BufferedStarts:       bufferedStarts,
		OngoingBackfills:     ongoingBackfills,
		LastCompletionResult: resultPayloads,
		ContinuedFailure:     continuedFailure,
		ConflictToken:        schedulerState.ConflictToken,
		NeedRefresh:          len(running) > 0,
	}

	return &schedulespb.StartScheduleArgs{
		Schedule: schedule,
		Info:     info,
		State:    state,
	}
}

// convertBufferedStartsLegacyToCHASM transforms V1 buffered starts to V2 format.
// V2 requires request_id, workflow_id, attempt, and backoff_time fields.
func convertBufferedStartsLegacyToCHASM(
	v1Starts []*schedulespb.BufferedStart,
	namespaceID, scheduleID string,
	conflictToken int64,
	baseWorkflowID string,
) []*schedulespb.BufferedStart {
	if len(v1Starts) == 0 {
		return nil
	}

	v2Starts := make([]*schedulespb.BufferedStart, len(v1Starts))
	for i, v1Start := range v1Starts {
		v2Start := common.CloneProto(v1Start)

		if v2Start.RequestId == "" {
			v2Start.RequestId = schedulescommon.GenerateRequestID(
				namespaceID,
				scheduleID,
				conflictToken,
				"migrated",
				v1Start.GetNominalTime().AsTime(),
				v1Start.GetActualTime().AsTime(),
			)
		}

		if v2Start.WorkflowId == "" {
			v2Start.WorkflowId = schedulescommon.GenerateWorkflowID(
				baseWorkflowID,
				v1Start.GetNominalTime().AsTime(),
			)
		}

		v2Start.Attempt = 0
		v2Start.BackoffTime = nil

		v2Starts[i] = v2Start
	}

	return v2Starts
}

// convertRunningWorkflowsToBufferedStarts converts V1's RunningWorkflows list to V2's
// BufferedStarts format. In V2, running workflows are represented as BufferedStarts with
// RunId and StartTime populated, and Completed field empty.
func convertRunningWorkflowsToBufferedStarts(
	runningWorkflows []*commonpb.WorkflowExecution,
	namespaceID, scheduleID string,
	conflictToken int64,
	migrationTime time.Time,
) []*schedulespb.BufferedStart {
	if len(runningWorkflows) == 0 {
		return nil
	}

	bufferedStarts := make([]*schedulespb.BufferedStart, len(runningWorkflows))
	for i, wf := range runningWorkflows {
		bufferedStarts[i] = &schedulespb.BufferedStart{
			NominalTime: timestamppb.New(migrationTime),
			ActualTime:  timestamppb.New(migrationTime),
			StartTime:   timestamppb.New(migrationTime),
			WorkflowId:  wf.WorkflowId,
			RunId:       wf.RunId,
			// RequestId will be used with AttachRequestID to register Nexus
			// callbacks for tracking workflow completion after migration.
			// Include the RunId in the tag to ensure each running workflow
			// gets a unique RequestId (important for ALLOW_ALL overlap
			// policy where multiple workflows may be running concurrently).
			RequestId: schedulescommon.GenerateRequestID(
				namespaceID,
				scheduleID,
				conflictToken,
				"migrated-running-"+wf.RunId,
				migrationTime,
				migrationTime,
			),
			Attempt:   1,
			Completed: nil,
			// Migrated running workflows must have a Nexus callback attached once the
			// migrated schedule target has been created.
			HasCallback: false,
		}
	}

	return bufferedStarts
}

// convertRecentActionsToBufferedStarts converts V1's RecentActions list to V2's
// BufferedStarts format. In V2, completed actions are represented as BufferedStarts with
// RunId, StartTime, and Completed fields all populated.
//
// runningWorkflows is the set of currently running workflow executions (from
// info.RunningWorkflows). These are excluded because they are already converted
// separately by convertRunningWorkflowsToBufferedStarts. In V1, recordAction
// adds the same workflow to both RecentActions and RunningWorkflows, so without
// this filter the same execution would appear twice in the CHASM BufferedStarts.
func convertRecentActionsToBufferedStarts(
	recentActions []*schedulepb.ScheduleActionResult,
	runningWorkflows []*commonpb.WorkflowExecution,
	namespaceID, scheduleID string,
	conflictToken int64,
	migrationTime time.Time,
) []*schedulespb.BufferedStart {
	if len(recentActions) == 0 {
		return nil
	}

	// Build a set of running workflow run IDs to exclude from recent actions,
	// since those are already converted by convertRunningWorkflowsToBufferedStarts.
	runningRunIDs := make(map[string]struct{}, len(runningWorkflows))
	for _, wf := range runningWorkflows {
		runningRunIDs[wf.GetRunId()] = struct{}{}
	}

	bufferedStarts := make([]*schedulespb.BufferedStart, 0, len(recentActions))
	for _, action := range recentActions {
		if action.StartWorkflowResult == nil {
			continue
		}

		// Skip actions for workflows that are still running — those are handled
		// by convertRunningWorkflowsToBufferedStarts.
		if _, ok := runningRunIDs[action.StartWorkflowResult.GetRunId()]; ok {
			continue
		}

		bufferedStarts = append(bufferedStarts, &schedulespb.BufferedStart{
			NominalTime: action.ScheduleTime,
			ActualTime:  action.ActualTime,
			StartTime:   action.ActualTime,
			WorkflowId:  action.StartWorkflowResult.WorkflowId,
			RunId:       action.StartWorkflowResult.RunId,
			RequestId: schedulescommon.GenerateRequestID(
				namespaceID,
				scheduleID,
				conflictToken,
				"migrated-completed",
				action.ScheduleTime.AsTime(),
				action.ActualTime.AsTime(),
			),
			Attempt: 1,
			Completed: &schedulespb.CompletedResult{
				Status:    action.StartWorkflowStatus,
				CloseTime: timestamppb.New(migrationTime),
			},
		})
	}

	return bufferedStarts
}

func convertBackfillsLegacyToCHASM(
	legacyBackfills []*schedulepb.BackfillRequest) map[string]*schedulerpb.BackfillerState {
	if len(legacyBackfills) == 0 {
		return nil
	}

	backfillers := make(map[string]*schedulerpb.BackfillerState, len(legacyBackfills))
	for _, v1Backfill := range legacyBackfills {
		backfillID := schedulescommon.GenerateBackfillerID()

		backfillers[backfillID] = &schedulerpb.BackfillerState{
			Request: &schedulerpb.BackfillerState_BackfillRequest{
				BackfillRequest: common.CloneProto(v1Backfill),
			},
			BackfillId:        backfillID,
			LastProcessedTime: nil,
			Attempt:           0,
		}
	}

	return backfillers
}

// convertLastCompletionLegacyToCHASM transforms V1 completion result to V2 format.
// V1 uses Payloads (plural), V2 uses single Payload.
func convertLastCompletionLegacyToCHASM(
	result *commonpb.Payloads,
	failure *failurepb.Failure,
) *schedulerpb.LastCompletionResult {
	if result == nil && failure == nil {
		return nil
	}

	lcr := &schedulerpb.LastCompletionResult{}
	if result != nil && len(result.Payloads) > 0 {
		lcr.Success = common.CloneProto(result.Payloads[0])
	}

	if failure != nil {
		lcr.Failure = common.CloneProto(failure)
	}

	return lcr
}

// getWorkflowID extracts the workflow ID from the schedule's action.
// This is the workflow ID specified in the schedule spec. During workflow start
// generation, nominal time is suffixed to this ID.
func getWorkflowID(schedule *schedulepb.Schedule) string {
	if schedule == nil {
		return ""
	}
	return schedule.GetAction().GetStartWorkflow().GetWorkflowId()
}

func splitBufferedStartsForLegacy(
	starts []*schedulespb.BufferedStart,
) ([]*schedulespb.BufferedStart, []*commonpb.WorkflowExecution, []*schedulepb.ScheduleActionResult) {
	if len(starts) == 0 {
		return nil, nil, nil
	}

	var buffered []*schedulespb.BufferedStart
	var running []*commonpb.WorkflowExecution
	var recent []*schedulepb.ScheduleActionResult

	for _, start := range starts {
		if start.GetRunId() == "" && start.GetCompleted() == nil {
			buffered = append(buffered, common.CloneProto(start))
			continue
		}

		if start.GetRunId() == "" {
			continue
		}

		status := enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING
		if start.GetCompleted() != nil {
			status = start.GetCompleted().GetStatus()
		}

		recent = append(recent, &schedulepb.ScheduleActionResult{
			ScheduleTime: start.GetActualTime(),
			ActualTime:   start.GetStartTime(),
			StartWorkflowResult: &commonpb.WorkflowExecution{
				WorkflowId: start.GetWorkflowId(),
				RunId:      start.GetRunId(),
			},
			StartWorkflowStatus: status,
		})

		if start.GetCompleted() == nil {
			running = append(running, &commonpb.WorkflowExecution{
				WorkflowId: start.GetWorkflowId(),
				RunId:      start.GetRunId(),
			})
		}
	}

	return buffered, running, recent
}

func convertBackfillersCHASMToLegacy(
	backfillers map[string]*schedulerpb.BackfillerState,
	migrationTime time.Time,
) ([]*schedulepb.BackfillRequest, []*schedulespb.BufferedStart) {
	if len(backfillers) == 0 {
		return nil, nil
	}

	var ongoing []*schedulepb.BackfillRequest
	var triggerStarts []*schedulespb.BufferedStart

	for _, backfiller := range backfillers {
		if request := backfiller.GetBackfillRequest(); request != nil {
			backfill := common.CloneProto(request)
			if backfiller.GetAttempt() > 0 && backfiller.GetLastProcessedTime() != nil {
				backfill.StartTime = common.CloneProto(backfiller.GetLastProcessedTime())
			}
			ongoing = append(ongoing, backfill)
			continue
		}

		if trigger := backfiller.GetTriggerRequest(); trigger != nil {
			when := backfiller.GetLastProcessedTime()
			if when == nil {
				when = timestamppb.New(migrationTime)
			}
			triggerStarts = append(triggerStarts, &schedulespb.BufferedStart{
				NominalTime:   when,
				ActualTime:    when,
				DesiredTime:   when,
				OverlapPolicy: trigger.GetOverlapPolicy(),
				Manual:        true,
			})
		}
	}

	return ongoing, triggerStarts
}

func convertLastCompletionCHASMToLegacy(
	result *schedulerpb.LastCompletionResult,
) (*commonpb.Payloads, *failurepb.Failure) {
	if result == nil {
		return nil, nil
	}

	var payloads *commonpb.Payloads
	if result.Success != nil {
		payloads = &commonpb.Payloads{
			Payloads: []*commonpb.Payload{common.CloneProto(result.Success)},
		}
	}

	return payloads, common.CloneProto(result.Failure)
}