geo/apps/desktop-client/src/main/publish-scheduler.ts

export type PublishSchedulerRouting = 'rabbitmq-primary' | 'db-recovery'

type PublishTaskState = 'queued' | 'leasing' | 'active'

export interface PublishSchedulerTaskRecord {
  taskId: string
  platform: string
  routing: PublishSchedulerRouting
  state: PublishTaskState
  enqueuedAt: number
  updatedAt: number
  lastStartedAt: number | null
}

export interface PublishSchedulerSelection {
  taskId: string
  platform: string
  routing: PublishSchedulerRouting
}

export interface PublishSchedulerSnapshot {
  queueDepth: number
  queuedCount: number
  leasingCount: number
  activeCount: number
  cooldownCount: number
  tasks: PublishSchedulerTaskRecord[]
  cooldownUntil: Record<string, number>
}

interface PublishSchedulerSelectionOptions {
  now?: number
  maxConcurrency: number
  activePlatforms?: ReadonlySet<string>
  activeCount: number
}

const publishPlatformCooldownMinMs = 3_000
const publishPlatformCooldownMaxMs = 6_000

const tasks = new Map<string, PublishSchedulerTaskRecord>()
const cooldownUntil = new Map<string, number>()

export function initPublishScheduler(): void {
  tasks.clear()
  cooldownUntil.clear()
}

export function enqueuePublishTask(input: {
  taskId: string
  platform: string
  routing: PublishSchedulerRouting
}): void {
  const taskId = input.taskId.trim()
  const platform = input.platform.trim()
  if (!taskId || !platform) {
    return
  }

  const now = Date.now()
  const existing = tasks.get(taskId)
  tasks.set(taskId, {
    taskId,
    platform,
    routing: input.routing,
    state: existing?.state === 'active' ? 'active' : 'queued',
    enqueuedAt: existing?.enqueuedAt ?? now,
    updatedAt: now,
    lastStartedAt: existing?.lastStartedAt ?? null,
  })
}

export function selectNextPublishTask(
  options: PublishSchedulerSelectionOptions,
): PublishSchedulerSelection | null {
  const now = options.now ?? Date.now()
  pruneCooldowns(now)

  if (options.activeCount >= options.maxConcurrency) {
    return null
  }

  const activePlatforms = new Set(options.activePlatforms ?? [])
  for (const task of tasks.values()) {
    if (task.state === 'active' || task.state === 'leasing') {
      activePlatforms.add(task.platform)
    }
  }

  const candidates = [...tasks.values()]
    .filter((task) => task.state === 'queued')
    .sort((left, right) => {
      if (left.enqueuedAt !== right.enqueuedAt) {
        return left.enqueuedAt - right.enqueuedAt
      }
      return left.updatedAt - right.updatedAt
    })

  for (const candidate of candidates) {
    if (activePlatforms.has(candidate.platform)) {
      continue
    }

    const platformCooldown = cooldownUntil.get(candidate.platform) ?? 0
    if (platformCooldown > now) {
      continue
    }

    candidate.state = 'leasing'
    candidate.updatedAt = now
    tasks.set(candidate.taskId, candidate)
    return {
      taskId: candidate.taskId,
      platform: candidate.platform,
      routing: candidate.routing,
    }
  }

  return null
}

export function notePublishTaskActivated(taskId: string, platform: string): void {
  const normalizedTaskId = taskId.trim()
  const normalizedPlatform = platform.trim()
  if (!normalizedTaskId || !normalizedPlatform) {
    return
  }

  const now = Date.now()
  const existing = tasks.get(normalizedTaskId)
  tasks.set(normalizedTaskId, {
    taskId: normalizedTaskId,
    platform: normalizedPlatform,
    routing: existing?.routing ?? 'db-recovery',
    state: 'active',
    enqueuedAt: existing?.enqueuedAt ?? now,
    updatedAt: now,
    lastStartedAt: now,
  })
}

export function notePublishTaskLeaseMiss(taskId: string): void {
  tasks.delete(taskId)
}

export function notePublishTaskCompleted(taskId: string, platform?: string): void {
  const existing = tasks.get(taskId)
  const resolvedPlatform = (platform ?? existing?.platform ?? '').trim()
  tasks.delete(taskId)

  if (resolvedPlatform) {
    cooldownUntil.set(resolvedPlatform, Date.now() + randomPublishCooldownMs())
  }
}

export function notePublishTaskCanceled(taskId: string): void {
  tasks.delete(taskId)
}

export function getPublishSchedulerSnapshot(): PublishSchedulerSnapshot {
  const now = Date.now()
  pruneCooldowns(now)
  const sortedTasks = [...tasks.values()].sort((left, right) => {
    if (left.state !== right.state) {
      return stateRank(left.state) - stateRank(right.state)
    }
    return left.enqueuedAt - right.enqueuedAt
  })

  return {
    queueDepth: sortedTasks.filter((task) => task.state === 'queued' || task.state === 'leasing')
      .length,
    queuedCount: sortedTasks.filter((task) => task.state === 'queued').length,
    leasingCount: sortedTasks.filter((task) => task.state === 'leasing').length,
    activeCount: sortedTasks.filter((task) => task.state === 'active').length,
    cooldownCount: cooldownUntil.size,
    tasks: sortedTasks.map((task) => ({ ...task })),
    cooldownUntil: Object.fromEntries(cooldownUntil.entries()),
  }
}

function pruneCooldowns(now: number): void {
  for (const [platform, until] of cooldownUntil.entries()) {
    if (until <= now) {
      cooldownUntil.delete(platform)
    }
  }
}

function randomPublishCooldownMs(): number {
  const spread = publishPlatformCooldownMaxMs - publishPlatformCooldownMinMs
  return publishPlatformCooldownMinMs + Math.floor(Math.random() * (spread + 1))
}

function stateRank(state: PublishTaskState): number {
  if (state === 'active') {
    return 0
  }
  if (state === 'leasing') {
    return 1
  }
  return 2
}