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 }