EvoBGP/internal/store/memory_job_queue.go

package store

import (
	"context"
	"encoding/json"
	"os"
	"sort"
	"strconv"
	"sync"
	"time"

	"github.com/google/uuid"
)

func copyJob(j JobRecord) JobRecord {
	if j.IdempotencyKey != nil {
		v := *j.IdempotencyKey
		j.IdempotencyKey = &v
	}
	if j.ModuleID != nil {
		v := *j.ModuleID
		j.ModuleID = &v
	}
	if j.Error != nil {
		v := *j.Error
		j.Error = &v
	}
	if j.StartedAt != nil {
		v := *j.StartedAt
		j.StartedAt = &v
	}
	if j.FinishedAt != nil {
		v := *j.FinishedAt
		j.FinishedAt = &v
	}
	if j.LeaseUntil != nil {
		v := *j.LeaseUntil
		j.LeaseUntil = &v
	}
	if j.Meta != nil {
		raw, _ := json.Marshal(j.Meta)
		var meta map[string]any
		_ = json.Unmarshal(raw, &meta)
		j.Meta = meta
	}
	return j
}

func (m *Memory) EnqueueJob(_ context.Context, in JobEnqueue) (JobRecord, bool, error) {
	h, err := JobRequestHash(in)
	if err != nil {
		return JobRecord{}, false, err
	}
	unlockTenant := m.lockJobEnqueueTenant(in.TenantID)
	defer unlockTenant()
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	if m.jobs == nil {
		m.jobs = make(map[string]JobRecord)
		m.jobHashes = make(map[string]string)
		m.jobRequests = make(map[string]map[string]any)
	}
	if in.IdempotencyKey != nil && *in.IdempotencyKey != "" {
		for _, j := range m.jobs {
			if j.TenantID == in.TenantID && j.IdempotencyKey != nil && *j.IdempotencyKey == *in.IdempotencyKey {
				if m.jobHashes[j.ID] != h {
					return JobRecord{}, false, ErrJobConflict
				}
				return copyJob(j), false, nil
			}
		}
	}
	global, tenant := 0, 0
	for _, j := range m.jobs {
		if j.Status == "queued" || j.Status == "running" {
			global++
			if j.TenantID == in.TenantID {
				tenant++
			}
		}
	}
	if global >= jobLimitMemory("EVOBGP_JOB_MAX_PENDING", 10000) || tenant >= jobLimitMemory("EVOBGP_JOB_MAX_TENANT_PENDING", 1000) {
		return JobRecord{}, false, ErrJobQueueFull
	}
	j := JobRecord{ID: uuid.NewString(), TenantID: in.TenantID, Kind: in.Kind, Status: "queued", IdempotencyKey: in.IdempotencyKey, ModuleID: in.ModuleID, Meta: in.Meta, CreatedAt: time.Now().UTC()}
	j = copyJob(j)
	m.jobSequence++
	j.RequestSequence = m.jobSequence
	j.RequestHash = h
	m.jobs[j.ID] = j
	m.jobHashes[j.ID] = h
	m.jobRequests[j.ID] = copyJob(j).Meta
	return copyJob(j), true, nil
}

func (m *Memory) ClaimJobs(_ context.Context, owner string, limit int, lease time.Duration) ([]JobRecord, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	if limit <= 0 {
		return nil, nil
	}
	if limit > 64 {
		limit = 64
	}
	if lease <= 0 {
		lease = time.Minute
	}
	now := time.Now().UTC()
	m.reapExpiredJobsLocked(now)
	var all []JobRecord
	for _, j := range m.jobs {
		if j.Attempt < 3 && (j.Status == "queued" || (j.Status == "running" && j.LeaseUntil != nil && now.After(*j.LeaseUntil))) {
			all = append(all, j)
		}
	}
	sort.Slice(all, func(i, j int) bool {
		if !all[i].CreatedAt.Equal(all[j].CreatedAt) {
			return all[i].CreatedAt.Before(all[j].CreatedAt)
		}
		if all[i].RequestSequence != all[j].RequestSequence {
			return all[i].RequestSequence < all[j].RequestSequence
		}
		return all[i].ID < all[j].ID
	})
	if len(all) > limit {
		all = all[:limit]
	}
	for i := range all {
		j := all[i]
		j.Status = "running"
		j.OwnerID = owner
		j.FencingToken++
		j.Attempt++
		if j.StartedAt == nil {
			j.StartedAt = &now
		}
		until := now.Add(lease)
		j.LeaseUntil = &until
		m.jobs[j.ID] = j
		all[i] = copyJob(j)
	}
	return all, nil
}

// ReapExpiredJobs marks orphaned cancellations and exhausted retries terminal.
func (m *Memory) ReapExpiredJobs(ctx context.Context) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	m.reapExpiredJobsLocked(time.Now().UTC())
	return nil
}

func (m *Memory) reapExpiredJobsLocked(now time.Time) {
	for id, j := range m.jobs {
		if j.Status == "running" && j.LeaseUntil != nil && now.After(*j.LeaseUntil) && j.CancelRequested {
			j.Status = "cancelled"
			j.FinishedAt = &now
			j.OwnerID = ""
			j.LeaseUntil = nil
			m.jobs[id] = j
			continue
		}
		if (j.Status == "queued" || (j.Status == "running" && j.LeaseUntil != nil && now.After(*j.LeaseUntil))) && j.Attempt >= 3 {
			j.Status = "failed"
			msg := "retry limit reached"
			j.Error = &msg
			j.FinishedAt = &now
			j.OwnerID = ""
			j.LeaseUntil = nil
			m.jobs[id] = j
		}
	}
}

func (m *Memory) HeartbeatJob(_ context.Context, id, owner string, token int64, lease time.Duration) (bool, error) {
	if lease <= 0 {
		lease = time.Minute
	}
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	j, ok := m.jobs[id]
	if !ok || j.OwnerID != owner || j.FencingToken != token || j.Status != "running" || j.LeaseUntil == nil || !time.Now().Before(*j.LeaseUntil) {
		return false, ErrJobLeaseLost
	}
	until := time.Now().Add(lease)
	j.LeaseUntil = &until
	m.jobs[id] = j
	return j.CancelRequested, nil
}

func (m *Memory) FinishJob(_ context.Context, j JobRecord) error {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	old, ok := m.jobs[j.ID]
	if !ok || old.OwnerID != j.OwnerID || old.FencingToken != j.FencingToken || old.Status != "running" || old.LeaseUntil == nil || !time.Now().Before(*old.LeaseUntil) {
		return ErrJobLeaseLost
	}
	j.CreatedAt = old.CreatedAt
	j.RequestSequence = old.RequestSequence
	j.RequestHash = old.RequestHash
	j.IdempotencyKey = old.IdempotencyKey
	j.ModuleID = old.ModuleID
	j.TenantID = old.TenantID
	j.Kind = old.Kind
	j.StartedAt = old.StartedAt
	j.Attempt = old.Attempt
	j.CancelRequested = old.CancelRequested
	j.OwnerID = ""
	j.LeaseUntil = nil
	now := time.Now().UTC()
	j.FinishedAt = &now
	if old.CancelRequested {
		j.Status = "cancelled"
		j.Error = nil
	}
	m.jobs[j.ID] = copyJob(j)
	return nil
}

func (m *Memory) CheckpointJob(_ context.Context, in JobRecord) error {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	j, ok := m.jobs[in.ID]
	if !ok || j.OwnerID != in.OwnerID || j.FencingToken != in.FencingToken || j.Status != "running" || j.LeaseUntil == nil || !time.Now().Before(*j.LeaseUntil) {
		return ErrJobLeaseLost
	}
	j.Meta = copyJob(in).Meta
	m.jobs[j.ID] = j
	return nil
}

func (m *Memory) ReleaseJob(_ context.Context, id, owner string, token int64) error {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	j, ok := m.jobs[id]
	if !ok || j.OwnerID != owner || j.FencingToken != token || j.Status != "running" {
		return ErrJobLeaseLost
	}
	j.Status = "queued"
	if j.CancelRequested {
		j.Status = "cancelled"
		now := time.Now().UTC()
		j.FinishedAt = &now
	}
	j.OwnerID = ""
	j.LeaseUntil = nil
	if j.Attempt > 0 {
		j.Attempt--
	}
	m.jobs[id] = j
	return nil
}

func (m *Memory) GetJob(_ context.Context, tenant, id string) (JobRecord, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	j, ok := m.jobs[id]
	if !ok || j.TenantID != tenant {
		return JobRecord{}, ErrJobNotFound
	}
	return copyJob(j), nil
}

func (m *Memory) ListJobs(_ context.Context, tenant, status, kind, cursor string, limit int) ([]JobRecord, string, bool, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	var all []JobRecord
	for _, j := range m.jobs {
		if j.TenantID == tenant && (status == "" || j.Status == status) && (kind == "" || j.Kind == kind) {
			all = append(all, copyJob(j))
		}
	}
	sort.Slice(all, func(i, j int) bool {
		if !all[i].CreatedAt.Equal(all[j].CreatedAt) {
			return all[i].CreatedAt.After(all[j].CreatedAt)
		}
		if all[i].RequestSequence != all[j].RequestSequence {
			return all[i].RequestSequence > all[j].RequestSequence
		}
		return all[i].ID > all[j].ID
	})
	off, _ := strconv.Atoi(cursor)
	if off < 0 {
		off = 0
	}
	if limit <= 0 || limit > 100 {
		limit = 50
	}
	if off >= len(all) {
		return nil, "", false, nil
	}
	end := off + limit
	more := end < len(all)
	if end > len(all) {
		end = len(all)
	}
	next := ""
	if more {
		next = strconv.Itoa(end)
	}
	return all[off:end], next, more, nil
}

func (m *Memory) CancelJob(_ context.Context, tenant, id string) (JobRecord, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	j, ok := m.jobs[id]
	if !ok || j.TenantID != tenant {
		return JobRecord{}, ErrJobNotFound
	}
	if j.Status == "queued" {
		now := time.Now().UTC()
		j.Status = "cancelled"
		j.FinishedAt = &now
	}
	if j.Status == "running" {
		j.CancelRequested = true
	}
	m.jobs[id] = j
	return copyJob(j), nil
}

func (m *Memory) RetryJob(_ context.Context, tenant, id string) (JobRecord, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	j, ok := m.jobs[id]
	if !ok || j.TenantID != tenant {
		return JobRecord{}, ErrJobNotFound
	}
	if j.Status != "failed" {
		return JobRecord{}, ErrJobConflict
	}
	global, pending := 0, 0
	for _, item := range m.jobs {
		if item.Status == "queued" || item.Status == "running" {
			global++
			if item.TenantID == tenant {
				pending++
			}
		}
	}
	if global >= jobLimitMemory("EVOBGP_JOB_MAX_PENDING", 10000) || pending >= jobLimitMemory("EVOBGP_JOB_MAX_TENANT_PENDING", 1000) {
		return JobRecord{}, ErrJobQueueFull
	}
	j.Status = "queued"
	j.Error = nil
	j.FinishedAt = nil
	j.CancelRequested = false
	j.Attempt = 0
	j.StartedAt = nil
	j.OwnerID = ""
	j.LeaseUntil = nil
	j.Meta = copyJob(JobRecord{Meta: m.jobRequests[id]}).Meta
	m.jobs[id] = j
	return copyJob(j), nil
}

// lockJobEnqueueTenant serializes idempotency and the tenant pending limit.
// The global jobMu is still required to mutate the shared job map.
func (m *Memory) lockJobEnqueueTenant(tenantID string) func() {
	m.jobTenantMu.Lock()
	if m.jobTenantLocks == nil {
		m.jobTenantLocks = make(map[string]*sync.Mutex)
	}
	mu := m.jobTenantLocks[tenantID]
	if mu == nil {
		mu = &sync.Mutex{}
		m.jobTenantLocks[tenantID] = mu
	}
	m.jobTenantMu.Unlock()
	mu.Lock()
	return mu.Unlock
}

func jobLimitMemory(name string, fallback int) int {
	n, e := strconv.Atoi(os.Getenv(name))
	if e == nil && n > 0 {
		return n
	}
	return fallback
}

func (m *Memory) CountJobsByStatus(_ context.Context) (map[string]int64, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	out := map[string]int64{}
	for _, j := range m.jobs {
		out[j.Status]++
	}
	return out, nil
}

func (m *Memory) FindJobByIdempotencyKey(_ context.Context, tenant, key string) (JobRecord, error) {
	m.jobMu.Lock()
	defer m.jobMu.Unlock()
	for _, j := range m.jobs {
		if j.TenantID == tenant && j.IdempotencyKey != nil && *j.IdempotencyKey == key {
			return copyJob(j), nil
		}
	}
	return JobRecord{}, ErrJobNotFound
}