148 lines
3.0 KiB
Go
148 lines
3.0 KiB
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 {
|
||
s.mu.Lock()
|
||
defer s.mu.Unlock()
|
||
if _, ok := s.byID[j.ID]; !ok {
|
||
return domain.ErrNotFound
|
||
}
|
||
cp := *j
|
||
s.byID[j.ID] = &cp
|
||
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) 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 {
|
||
if j.Status != domain.StatusPending && j.Status != domain.StatusQueued {
|
||
continue
|
||
}
|
||
if j.RunAfter > 0 && j.RunAfter > now {
|
||
continue
|
||
}
|
||
if best == nil || j.CreatedAt < best.CreatedAt {
|
||
cp := *j
|
||
best = &cp
|
||
}
|
||
}
|
||
if best == nil {
|
||
return nil, domain.ErrNotFound
|
||
}
|
||
from := best.Status
|
||
to := domain.StatusRunning
|
||
if !domain.CanTransition(from, to) {
|
||
return nil, domain.ErrIllegalStatus
|
||
}
|
||
best.Status = to
|
||
best.WorkerID = workerID
|
||
best.UpdatedAt = now
|
||
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
|
||
}
|
||
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)
|