thread-master/apps/backend/internal/module/job/repository/memory.go

package repository

import (
	"context"
	"sync"

	"apps/backend/internal/module/job/domain"
)

type MemoryStore struct {
	mu   sync.Mutex
	byID map[string]*domain.Job
}

func NewMemory() *MemoryStore {
	return &MemoryStore{byID: map[string]*domain.Job{}}
}

func (s *MemoryStore) Insert(_ context.Context, j *domain.Job) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	cp := *j
	s.byID[j.ID] = &cp
	return nil
}

func (s *MemoryStore) Update(_ context.Context, j *domain.Job) error {
	return s.update(j, "")
}

func (s *MemoryStore) UpdateOwned(_ context.Context, j *domain.Job, leaseOwner string) error {
	return s.update(j, leaseOwner)
}

func (s *MemoryStore) update(j *domain.Job, leaseOwner string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	stored, ok := s.byID[j.ID]
	if !ok {
		return domain.ErrNotFound
	}
	if stored.Version != j.Version {
		return domain.ErrIllegalStatus
	}
	if leaseOwner != "" && (stored.Status != domain.StatusRunning || stored.LeaseOwner != leaseOwner) {
		return domain.ErrIllegalStatus
	}
	cp := *j
	cp.Version++
	s.byID[j.ID] = &cp
	j.Version = cp.Version
	return nil
}

func (s *MemoryStore) RenewLease(_ context.Context, id, leaseOwner string, leaseExpiresAt int64) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.byID[id]
	if !ok {
		return domain.ErrNotFound
	}
	if j.Status != domain.StatusRunning || j.LeaseOwner != leaseOwner {
		return domain.ErrIllegalStatus
	}
	j.LeaseExpiresAt = leaseExpiresAt
	j.UpdatedAt = domain.NowNano()
	return nil
}

func (s *MemoryStore) FindByID(_ context.Context, id string) (*domain.Job, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	j, ok := s.byID[id]
	if !ok {
		return nil, domain.ErrNotFound
	}
	cp := *j
	return &cp, nil
}

func (s *MemoryStore) ListByOwner(_ context.Context, ownerUID int64) ([]*domain.Job, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var out []*domain.Job
	for _, j := range s.byID {
		if j.OwnerUID == ownerUID {
			cp := *j
			out = append(out, &cp)
		}
	}
	return out, nil
}

func (s *MemoryStore) ListByOwnerTab(_ context.Context, ownerUID int64, tab string, page, pageSize int) ([]*domain.Job, int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	page, pageSize = domain.ClampPage(page, pageSize)
	now := domain.NowNano()
	var matched []*domain.Job
	for _, j := range s.byID {
		if j.OwnerUID != ownerUID {
			continue
		}
		if domain.MatchListTab(j, tab, now) {
			cp := *j
			matched = append(matched, &cp)
		}
	}
	// sort updated_at desc（recurring: run_after asc）
	tabN := domain.NormalizeListTab(tab)
	for i := 0; i < len(matched); i++ {
		for k := i + 1; k < len(matched); k++ {
			swap := false
			if tabN == domain.TabRecurring {
				if matched[k].RunAfter < matched[i].RunAfter {
					swap = true
				}
			} else if matched[k].UpdatedAt > matched[i].UpdatedAt {
				swap = true
			}
			if swap {
				matched[i], matched[k] = matched[k], matched[i]
			}
		}
	}
	total := int64(len(matched))
	start := (page - 1) * pageSize
	if start >= len(matched) {
		return []*domain.Job{}, total, nil
	}
	end := start + pageSize
	if end > len(matched) {
		end = len(matched)
	}
	return matched[start:end], total, nil
}

func (s *MemoryStore) ClaimNext(_ context.Context, workerID string) (*domain.Job, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	now := domain.NowNano()
	var best *domain.Job
	for _, j := range s.byID {
		queuedDue := (j.Status == domain.StatusPending || j.Status == domain.StatusQueued) && (j.RunAfter <= 0 || j.RunAfter <= now)
		staleRunning := j.Status == domain.StatusRunning && j.LeaseExpiresAt <= now
		if !queuedDue && !staleRunning {
			continue
		}
		if best == nil || j.CreatedAt < best.CreatedAt {
			cp := *j
			best = &cp
		}
	}
	if best == nil {
		return nil, domain.ErrNotFound
	}
	best.Status = domain.StatusRunning
	best.WorkerID = workerID
	best.LeaseOwner = workerID
	best.LeaseExpiresAt = now + domain.DefaultLeaseDurationNs
	best.Attempt++
	best.UpdatedAt = now
	best.Version++
	if best.ProgressSummary == "" {
		best.ProgressSummary = "已由 worker 領取（" + workerID + "）"
	}
	s.byID[best.ID] = best
	cp := *best
	return &cp, nil
}

func (s *MemoryStore) CancelPendingByRef(_ context.Context, ownerUID int64, template, refID string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	now := domain.NowNano()
	for _, j := range s.byID {
		if j.OwnerUID != ownerUID || j.TemplateType != template || j.RefID != refID {
			continue
		}
		if j.Status != domain.StatusPending && j.Status != domain.StatusQueued {
			continue
		}
		j.Status = domain.StatusCancelled
		j.ProgressSummary = "已由新的定期排程取代"
		j.CompletedAt = now
		j.UpdatedAt = now
		j.Version++
	}
	return nil
}

func (s *MemoryStore) Delete(_ context.Context, id string) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if _, ok := s.byID[id]; !ok {
		return domain.ErrNotFound
	}
	delete(s.byID, id)
	return nil
}

func (s *MemoryStore) DeleteTerminalBefore(_ context.Context, beforeNs int64) (int64, error) {
	s.mu.Lock()
	defer s.mu.Unlock()
	var n int64
	for id, j := range s.byID {
		if !domain.IsTerminal(j.Status) {
			continue
		}
		doneAt := j.CompletedAt
		if doneAt <= 0 {
			doneAt = j.UpdatedAt
		}
		if doneAt > 0 && doneAt < beforeNs {
			delete(s.byID, id)
			n++
		}
	}
	return n, nil
}

var _ domain.Repository = (*MemoryStore)(nil)