// Package localstore provides process-local development persistence. // It is intentionally non-durable and must not be used as a production store. package localstore import ( "context" "encoding/json" "fmt" "reflect" "sort" "sync" "time" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/administration" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/assets" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/billing" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/jobs" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/templates" "git.nianxx.cn/wangxuming/NianAIGC/backend/internal/usage" ) const ( DemoAccountID = "demo-merchant" DemoOrganizationID = "org-demo" DemoPhone = "13800000000" DemoPassword = "demo-password" ) type Option func(*Store) func WithClock(now func() time.Time) Option { return func(s *Store) { if now != nil { s.now = now } } } type Store struct { mu sync.RWMutex now func() time.Time accounts map[string]administration.Account organizations map[string]administration.Organization assets map[string]assets.Asset jobs map[string]jobs.Job templates map[string]templates.Template usageEvents map[string]usage.Event usageJobIDs map[string]string priceRules map[string]billing.PriceRule wallets map[string]billing.Wallet ledger []billing.LedgerEntry postings map[string]walletRecord } type walletRecord struct { params billing.WalletPostParams posting billing.WalletPosting } func New(options ...Option) *Store { s := &Store{now: time.Now, accounts: map[string]administration.Account{}, organizations: map[string]administration.Organization{}, assets: map[string]assets.Asset{}, jobs: map[string]jobs.Job{}, templates: map[string]templates.Template{}, usageEvents: map[string]usage.Event{}, usageJobIDs: map[string]string{}, priceRules: map[string]billing.PriceRule{}, wallets: map[string]billing.Wallet{}, postings: map[string]walletRecord{}} for _, option := range options { option(s) } s.seedDemo() return s } func NewStore(options ...Option) *Store { return New(options...) } func (s *Store) seedDemo() { now := s.now().UTC() password, err := administration.HashPassword(DemoPassword) if err != nil { panic(fmt.Sprintf("localstore: hash demo password: %v", err)) } s.organizations[DemoOrganizationID] = administration.Organization{ID: DemoOrganizationID, Name: "演示组织", Status: administration.StatusActive, ArchiveOwnerID: DemoAccountID, CreatedAt: now, UpdatedAt: now} s.accounts[DemoAccountID] = administration.Account{ID: DemoAccountID, Phone: DemoPhone, DisplayName: "智念用户", Role: administration.RoleSuperAdmin, OrganizationID: DemoOrganizationID, Status: administration.StatusActive, PasswordHash: password.Hash, PasswordSalt: password.Salt, SessionVersion: 1, CreatedAt: now, UpdatedAt: now} } func cloneAsset(a assets.Asset) assets.Asset { a.Tags = append([]string(nil), a.Tags...) a.Metadata = cloneMap(a.Metadata) return a } func cloneJob(j jobs.Job) jobs.Job { j.InputAssetIDs = append([]string(nil), j.InputAssetIDs...) j.InputURLs = append([]string(nil), j.InputURLs...) j.OutputAssetIDs = append([]string(nil), j.OutputAssetIDs...) j.RequestPayload = append(json.RawMessage(nil), j.RequestPayload...) j.ResponsePayload = append(json.RawMessage(nil), j.ResponsePayload...) j.WebhookLastStatus = append(json.RawMessage(nil), j.WebhookLastStatus...) j.UsageContext = append(json.RawMessage(nil), j.UsageContext...) j.Billing = append(json.RawMessage(nil), j.Billing...) if j.Error != nil { e := *j.Error j.Error = &e } if j.ProviderDispatchStartedAt != nil { value := *j.ProviderDispatchStartedAt j.ProviderDispatchStartedAt = &value } if j.DispatchReadyAt != nil { value := *j.DispatchReadyAt j.DispatchReadyAt = &value } if j.FinalizedAt != nil { value := *j.FinalizedAt j.FinalizedAt = &value } return j } func cloneMap(m map[string]any) map[string]any { if m == nil { return nil } raw, err := json.Marshal(m) if err != nil { out := make(map[string]any, len(m)) for k, v := range m { out[k] = v } return out } var out map[string]any _ = json.Unmarshal(raw, &out) return out } func cloneRule(r billing.PriceRule) billing.PriceRule { r.Conditions = billing.Conditions(cloneMap(map[string]any(r.Conditions))) r.Source = cloneMap(r.Source) raw, _ := json.Marshal(r.Dimensions) _ = json.Unmarshal(raw, &r.Dimensions) return r } // Assets. func (s *Store) ListOwner(_ context.Context, owner string) ([]assets.Asset, error) { s.mu.RLock() defer s.mu.RUnlock() out := []assets.Asset{} for _, a := range s.assets { if a.OwnerID == owner { out = append(out, cloneAsset(a)) } } sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) }) return out, nil } func (s *Store) GetOwner(_ context.Context, owner, id string) (assets.Asset, bool, error) { s.mu.RLock() defer s.mu.RUnlock() a, ok := s.assets[id] if !ok || a.OwnerID != owner { return assets.Asset{}, false, nil } return cloneAsset(a), true, nil } func (s *Store) GetOwnerByStoragePath(_ context.Context, owner, path string) (assets.Asset, bool, error) { s.mu.RLock() defer s.mu.RUnlock() for _, a := range s.assets { if a.OwnerID == owner && a.StoragePath == path { return cloneAsset(a), true, nil } } return assets.Asset{}, false, nil } func (s *Store) ListPublic(_ context.Context, owner, client string, limit int) ([]assets.Asset, error) { s.mu.RLock() defer s.mu.RUnlock() allowed := s.publicAssetIDs(owner, client, limit) tag := assets.ClientTag(client) out := []assets.Asset{} for _, a := range s.assets { if a.OwnerID == owner && (contains(a.Tags, tag) || allowed[a.ID]) { out = append(out, cloneAsset(a)) } } sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) }) return out, nil } func (s *Store) GetPublic(_ context.Context, owner, client, id string, limit int) (assets.Asset, bool, error) { s.mu.RLock() defer s.mu.RUnlock() a, ok := s.assets[id] if !ok || a.OwnerID != owner { return assets.Asset{}, false, nil } if !contains(a.Tags, assets.ClientTag(client)) && !s.publicAssetIDs(owner, client, limit)[id] { return assets.Asset{}, false, nil } return cloneAsset(a), true, nil } func (s *Store) publicAssetIDs(owner, client string, limit int) map[string]bool { candidates := []jobs.Job{} for _, j := range s.jobs { if j.OwnerID == owner && j.ExternalClientID == client { candidates = append(candidates, j) } } sort.Slice(candidates, func(i, j int) bool { return candidates[i].CreatedAt.After(candidates[j].CreatedAt) }) if limit <= 0 { limit = 200 } if len(candidates) > limit { candidates = candidates[:limit] } out := map[string]bool{} for _, j := range candidates { for _, id := range append(append([]string{}, j.InputAssetIDs...), j.OutputAssetIDs...) { out[id] = true } } return out } func (s *Store) Create(_ context.Context, a assets.Asset) (assets.Asset, error) { s.mu.Lock() defer s.mu.Unlock() if _, ok := s.assets[a.ID]; ok { return assets.Asset{}, fmt.Errorf("asset already exists: %s", a.ID) } if a.Tags == nil { a.Tags = []string{} } if a.Metadata == nil { a.Metadata = map[string]any{} } s.assets[a.ID] = cloneAsset(a) return cloneAsset(a), nil } func (s *Store) DeleteOwner(_ context.Context, owner, id string) (assets.Asset, bool, error) { s.mu.Lock() defer s.mu.Unlock() a, ok := s.assets[id] if !ok || a.OwnerID != owner { return assets.Asset{}, false, nil } delete(s.assets, id) return cloneAsset(a), true, nil } // Jobs. func (s *Store) ListJobs(_ context.Context, f jobs.ListFilter) ([]jobs.Job, error) { s.mu.RLock() defer s.mu.RUnlock() out := []jobs.Job{} for _, j := range s.jobs { if !f.Scope.Owns(j) || f.Status != "" && j.Status != f.Status || f.Capability != "" && j.Capability != f.Capability || f.Before != nil && !j.CreatedAt.Before(*f.Before) { continue } out = append(out, cloneJob(j)) } sort.Slice(out, func(i, j int) bool { return out[i].CreatedAt.After(out[j].CreatedAt) }) if f.Limit > 0 && len(out) > f.Limit { out = out[:f.Limit] } return out, nil } func (s *Store) FindJob(_ context.Context, scope jobs.Scope, id string) (jobs.Job, bool, error) { s.mu.RLock() defer s.mu.RUnlock() j, ok := s.jobs[id] if !ok || !scope.Owns(j) { return jobs.Job{}, false, nil } return cloneJob(j), true, nil } func (s *Store) FindIdempotentJob(_ context.Context, scope jobs.Scope, key string) (jobs.Job, bool, error) { s.mu.RLock() defer s.mu.RUnlock() for _, j := range s.jobs { if scope.Owns(j) && j.IdempotencyKey == key { return cloneJob(j), true, nil } } return jobs.Job{}, false, nil } func (s *Store) CreateJob(_ context.Context, j jobs.Job) (jobs.Job, error) { s.mu.Lock() defer s.mu.Unlock() if _, ok := s.jobs[j.ID]; ok { return jobs.Job{}, fmt.Errorf("generation job already exists: %s", j.ID) } if j.ExternalClientID != "" && j.IdempotencyKey != "" { for _, v := range s.jobs { if v.OwnerID == j.OwnerID && v.ExternalClientID == j.ExternalClientID && v.IdempotencyKey == j.IdempotencyKey { return jobs.Job{}, jobs.ErrUniqueIdempotency } } } s.jobs[j.ID] = cloneJob(j) return cloneJob(j), nil } func (s *Store) UpdateJob(_ context.Context, id string, p jobs.Patch) (jobs.Job, error) { s.mu.Lock() defer s.mu.Unlock() j, ok := s.jobs[id] if !ok { return jobs.Job{}, fmt.Errorf("generation job not found: %s", id) } if len(p.ExpectedStatuses) > 0 && !containsStatus(p.ExpectedStatuses, j.Status) { return jobs.Job{}, jobs.ErrTransitionConflict } if p.ExpectedLockedBy != nil && j.LockedBy != *p.ExpectedLockedBy { return jobs.Job{}, jobs.ErrTransitionConflict } if p.Status != nil { j.Status = *p.Status } if p.Error != nil { e := *p.Error j.Error = &e } if p.ClearError { j.Error = nil } if p.Attempts != nil { j.Attempts = *p.Attempts } if p.ScheduledAt != nil { j.ScheduledAt = *p.ScheduledAt } if p.CompletedAt != nil { v := *p.CompletedAt j.CompletedAt = &v } if p.ProviderTaskID != nil { j.ProviderTaskID = *p.ProviderTaskID } if p.ProviderDispatchStartedAt != nil { started := *p.ProviderDispatchStartedAt j.ProviderDispatchStartedAt = &started } if p.ClearProviderDispatch { j.ProviderDispatchStartedAt = nil } if p.SetResponsePayload { j.ResponsePayload = append([]byte(nil), p.ResponsePayload...) } if p.DispatchReadyAt != nil { ready := *p.DispatchReadyAt j.DispatchReadyAt = &ready } if p.FinalizedAt != nil { finalized := *p.FinalizedAt j.FinalizedAt = &finalized } if p.ClearProviderTaskID { j.ProviderTaskID = "" } if p.ClearLease { j.LockedAt = nil j.LockedBy = "" } if p.WebhookAttempts != nil { j.WebhookAttempts = *p.WebhookAttempts } if p.SetWebhookStatus { j.WebhookLastStatus = append([]byte(nil), p.WebhookLastStatus...) } j.UpdatedAt = s.now().UTC() s.jobs[id] = cloneJob(j) return cloneJob(j), nil } func (s *Store) DeleteJob(_ context.Context, id string) error { s.mu.Lock() defer s.mu.Unlock() if _, ok := s.jobs[id]; !ok { return fmt.Errorf("generation job not found: %s", id) } delete(s.jobs, id) return nil } func (s *Store) ClaimJobs(_ context.Context, worker string, limit, timeoutSeconds int) ([]jobs.Job, error) { s.mu.Lock() defer s.mu.Unlock() if limit < 1 { limit = 1 } if limit > 20 { limit = 20 } if timeoutSeconds <= 0 { timeoutSeconds = 300 } now := s.now().UTC() eligible := []jobs.Job{} for _, j := range s.jobs { if j.DispatchReadyAt == nil || j.FinalizedAt != nil || j.ScheduledAt.After(now) { continue } leaseAvailable := j.LockedAt == nil || j.LockedAt.Add(time.Duration(timeoutSeconds)*time.Second).Before(now) if leaseAvailable && (j.Status == jobs.StatusQueued || j.Status == jobs.StatusRunning || j.Status.Terminal()) { eligible = append(eligible, j) } } sort.Slice(eligible, func(i, j int) bool { if eligible[i].Priority != eligible[j].Priority { return eligible[i].Priority > eligible[j].Priority } if !eligible[i].ScheduledAt.Equal(eligible[j].ScheduledAt) { return eligible[i].ScheduledAt.Before(eligible[j].ScheduledAt) } return eligible[i].CreatedAt.Before(eligible[j].CreatedAt) }) if len(eligible) > limit { eligible = eligible[:limit] } out := make([]jobs.Job, 0, len(eligible)) for _, j := range eligible { j.LockedBy = worker j.LockedAt = timePtr(now) if j.StartedAt == nil { j.StartedAt = timePtr(now) } j.UpdatedAt = now s.jobs[j.ID] = cloneJob(j) out = append(out, cloneJob(j)) } return out, nil } func (s *Store) WriteBilling(_ context.Context, id string, value json.RawMessage) error { if len(value) == 0 || !json.Valid(value) { return fmt.Errorf("write generation billing: invalid JSON") } s.mu.Lock() defer s.mu.Unlock() j, ok := s.jobs[id] if !ok { return fmt.Errorf("write generation billing: generation job not found") } j.Billing = append([]byte(nil), value...) j.UpdatedAt = s.now().UTC() s.jobs[id] = j return nil } func (s *Store) WriteBillingFenced(_ context.Context, id string, value json.RawMessage, status jobs.Status, lockedBy string) error { if len(value) == 0 || !json.Valid(value) { return fmt.Errorf("write generation billing: invalid JSON") } s.mu.Lock() defer s.mu.Unlock() j, ok := s.jobs[id] if !ok || j.Status != status || j.LockedBy != lockedBy { return jobs.ErrTransitionConflict } j.Billing = append([]byte(nil), value...) j.UpdatedAt = s.now().UTC() s.jobs[id] = cloneJob(j) return nil } func (s *Store) ActivateCreation(_ context.Context, id string, value json.RawMessage) error { if len(value) == 0 || !json.Valid(value) { return fmt.Errorf("activate generation creation: invalid JSON") } s.mu.Lock() defer s.mu.Unlock() j, ok := s.jobs[id] if !ok { return fmt.Errorf("activate generation creation: generation job not found") } if j.DispatchReadyAt != nil { if !jsonEqual(j.Billing, value) { return jobs.ErrTransitionConflict } return nil } j.Billing = append([]byte(nil), value...) ready := s.now().UTC() j.DispatchReadyAt = &ready j.UpdatedAt = s.now().UTC() s.jobs[id] = cloneJob(j) return nil } func jsonEqual(a, b json.RawMessage) bool { var left, right any return json.Unmarshal(a, &left) == nil && json.Unmarshal(b, &right) == nil && reflect.DeepEqual(left, right) } func (s *Store) WriteOutputAssetIDs(_ context.Context, id string, values []string) error { s.mu.Lock() defer s.mu.Unlock() j, ok := s.jobs[id] if !ok { return fmt.Errorf("write generation output assets: generation job not found") } j.OutputAssetIDs = append([]string{}, values...) j.UpdatedAt = s.now().UTC() s.jobs[id] = j return nil } func (s *Store) WriteOutputAssetIDsFenced(_ context.Context, id string, values []string, status jobs.Status, lockedBy string) error { s.mu.Lock() defer s.mu.Unlock() j, ok := s.jobs[id] if !ok || j.Status != status || j.LockedBy != lockedBy { return jobs.ErrTransitionConflict } j.OutputAssetIDs = append([]string{}, values...) j.UpdatedAt = s.now().UTC() s.jobs[id] = cloneJob(j) return nil } func (s *Store) FailCreation(_ context.Context, j jobs.Job) error { if j.Status != jobs.StatusFailed || j.Error == nil || len(j.Billing) == 0 || !json.Valid(j.Billing) { return fmt.Errorf("fail generation creation: invalid terminal state") } s.mu.Lock() defer s.mu.Unlock() current, ok := s.jobs[j.ID] if !ok { return fmt.Errorf("fail generation creation: generation job not found") } current.Status = j.Status e := *j.Error current.Error = &e current.Billing = append([]byte(nil), j.Billing...) now := s.now().UTC() if current.CompletedAt == nil { current.CompletedAt = &now } current.LockedAt = nil current.LockedBy = "" current.UpdatedAt = now s.jobs[j.ID] = current return nil } func timePtr(v time.Time) *time.Time { return &v } func contains(values []string, want string) bool { for _, v := range values { if v == want { return true } } return false } func containsStatus(values []jobs.Status, want jobs.Status) bool { for _, value := range values { if value == want { return true } } return false }