284 lines
8.2 KiB
Go
284 lines
8.2 KiB
Go
package repository
|
||
|
||
import (
|
||
"context"
|
||
"errors"
|
||
|
||
libmongo "apps/backend/internal/lib/mongo"
|
||
"apps/backend/internal/module/job/domain"
|
||
|
||
"github.com/zeromicro/go-zero/core/logx"
|
||
"github.com/zeromicro/go-zero/core/stores/mon"
|
||
"go.mongodb.org/mongo-driver/bson"
|
||
"go.mongodb.org/mongo-driver/mongo"
|
||
"go.mongodb.org/mongo-driver/mongo/options"
|
||
)
|
||
|
||
const colJobs = "jobs"
|
||
|
||
type MonStore struct {
|
||
jobs *mon.Model
|
||
claimJobs *mongo.Collection
|
||
}
|
||
|
||
func NewMonStore(uri, database string) *MonStore {
|
||
uri = libmongo.MustMongoURI(uri)
|
||
client, err := mongo.Connect(context.Background(), options.Client().ApplyURI(uri))
|
||
if err != nil {
|
||
panic(err)
|
||
}
|
||
return &MonStore{
|
||
jobs: mon.MustNewModel(uri, database, colJobs),
|
||
claimJobs: client.Database(database).Collection(colJobs),
|
||
}
|
||
}
|
||
|
||
func (s *MonStore) Insert(ctx context.Context, j *domain.Job) error {
|
||
_, err := s.jobs.InsertOne(ctx, j)
|
||
return err
|
||
}
|
||
|
||
func (s *MonStore) Update(ctx context.Context, j *domain.Job) error {
|
||
return s.update(ctx, j, nil)
|
||
}
|
||
|
||
func (s *MonStore) UpdateOwned(ctx context.Context, j *domain.Job, leaseOwner string) error {
|
||
filter := bson.M{"status": domain.StatusRunning, "lease_owner": leaseOwner}
|
||
return s.update(ctx, j, filter)
|
||
}
|
||
|
||
func (s *MonStore) update(ctx context.Context, j *domain.Job, guard bson.M) error {
|
||
if j == nil {
|
||
return domain.ErrNotFound
|
||
}
|
||
expectedVersion := j.Version
|
||
j.UpdatedAt = domain.NowNano()
|
||
candidate := *j
|
||
candidate.Version = expectedVersion + 1
|
||
filter := bson.M{"_id": j.ID, "version": expectedVersion}
|
||
for key, value := range guard {
|
||
filter[key] = value
|
||
}
|
||
res, err := s.jobs.ReplaceOne(ctx, filter, &candidate)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if res.MatchedCount == 0 {
|
||
return domain.ErrIllegalStatus
|
||
}
|
||
j.Version = candidate.Version
|
||
return nil
|
||
}
|
||
|
||
func (s *MonStore) RenewLease(ctx context.Context, id, leaseOwner string, leaseExpiresAt int64) error {
|
||
res, err := s.jobs.UpdateOne(ctx, bson.M{
|
||
"_id": id, "status": domain.StatusRunning, "lease_owner": leaseOwner,
|
||
}, bson.M{
|
||
"$set": bson.M{"lease_expires_at": leaseExpiresAt, "updated_at": domain.NowNano()},
|
||
})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if res.MatchedCount == 0 {
|
||
return domain.ErrIllegalStatus
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *MonStore) FindByID(ctx context.Context, id string) (*domain.Job, error) {
|
||
var j domain.Job
|
||
err := s.jobs.FindOne(ctx, &j, bson.M{"_id": id})
|
||
if err != nil {
|
||
if err == mon.ErrNotFound {
|
||
return nil, domain.ErrNotFound
|
||
}
|
||
return nil, err
|
||
}
|
||
return &j, nil
|
||
}
|
||
|
||
func (s *MonStore) ListByOwner(ctx context.Context, ownerUID int64) ([]*domain.Job, error) {
|
||
var list []*domain.Job
|
||
err := s.jobs.Find(ctx, &list, bson.M{"owner_uid": ownerUID},
|
||
options.Find().SetSort(bson.D{{Key: "updated_at", Value: -1}}))
|
||
return list, err
|
||
}
|
||
|
||
func jobTabFilter(ownerUID int64, tab string, nowNs int64) bson.M {
|
||
base := bson.M{"owner_uid": ownerUID}
|
||
switch domain.NormalizeListTab(tab) {
|
||
case domain.TabHistory:
|
||
base["status"] = bson.M{"$in": []string{
|
||
domain.StatusSucceeded, domain.StatusFailed, domain.StatusCancelled,
|
||
}}
|
||
case domain.TabRecurring:
|
||
base["status"] = bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}}
|
||
base["run_after"] = bson.M{"$gt": nowNs}
|
||
default: // active
|
||
base["$or"] = []bson.M{
|
||
{"status": domain.StatusRunning},
|
||
{
|
||
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
|
||
"$or": []bson.M{
|
||
{"run_after": bson.M{"$exists": false}},
|
||
{"run_after": 0},
|
||
{"run_after": bson.M{"$lte": nowNs}},
|
||
},
|
||
},
|
||
}
|
||
}
|
||
return base
|
||
}
|
||
|
||
func (s *MonStore) ListByOwnerTab(ctx context.Context, ownerUID int64, tab string, page, pageSize int) ([]*domain.Job, int64, error) {
|
||
page, pageSize = domain.ClampPage(page, pageSize)
|
||
now := domain.NowNano()
|
||
filter := jobTabFilter(ownerUID, tab, now)
|
||
|
||
total, err := s.jobs.CountDocuments(ctx, filter)
|
||
if err != nil {
|
||
return nil, 0, err
|
||
}
|
||
|
||
sort := bson.D{{Key: "updated_at", Value: -1}}
|
||
if domain.NormalizeListTab(tab) == domain.TabRecurring {
|
||
sort = bson.D{{Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}
|
||
} else if domain.NormalizeListTab(tab) == domain.TabHistory {
|
||
sort = bson.D{{Key: "completed_at", Value: -1}, {Key: "updated_at", Value: -1}}
|
||
}
|
||
|
||
skip := int64((page - 1) * pageSize)
|
||
var list []*domain.Job
|
||
err = s.jobs.Find(ctx, &list, filter,
|
||
options.Find().SetSort(sort).SetSkip(skip).SetLimit(int64(pageSize)))
|
||
if err != nil {
|
||
return nil, 0, err
|
||
}
|
||
return list, total, nil
|
||
}
|
||
|
||
func (s *MonStore) ClaimNext(ctx context.Context, workerID string) (*domain.Job, error) {
|
||
now := domain.NowNano()
|
||
filter := bson.M{
|
||
"$or": []bson.M{
|
||
{
|
||
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
|
||
"$or": []bson.M{
|
||
{"run_after": bson.M{"$exists": false}},
|
||
{"run_after": 0},
|
||
{"run_after": bson.M{"$lte": now}},
|
||
},
|
||
},
|
||
{
|
||
"status": domain.StatusRunning,
|
||
"$or": []bson.M{
|
||
{"lease_expires_at": bson.M{"$exists": false}},
|
||
{"lease_expires_at": 0},
|
||
{"lease_expires_at": bson.M{"$lte": now}},
|
||
},
|
||
},
|
||
},
|
||
}
|
||
update := bson.M{"$set": bson.M{
|
||
"status": domain.StatusRunning,
|
||
"worker_id": workerID,
|
||
"lease_owner": workerID,
|
||
"lease_expires_at": now + domain.DefaultLeaseDurationNs,
|
||
"updated_at": now,
|
||
"progress_summary": "已由 worker 領取（" + workerID + "）",
|
||
}, "$inc": bson.M{"version": 1, "attempt": 1}}
|
||
opts := options.FindOneAndUpdate().
|
||
SetSort(bson.D{{Key: "run_after", Value: 1}, {Key: "created_at", Value: 1}}).
|
||
SetReturnDocument(options.After)
|
||
var j domain.Job
|
||
err := s.claimJobs.FindOneAndUpdate(ctx, filter, update, opts).Decode(&j)
|
||
if err != nil {
|
||
if errors.Is(err, mongo.ErrNoDocuments) {
|
||
return nil, domain.ErrNotFound
|
||
}
|
||
return nil, err
|
||
}
|
||
return &j, nil
|
||
}
|
||
|
||
func (s *MonStore) CancelPendingByRef(ctx context.Context, ownerUID int64, template, refID string) error {
|
||
if refID == "" {
|
||
return nil
|
||
}
|
||
filter := bson.M{
|
||
"owner_uid": ownerUID,
|
||
"template_type": template,
|
||
"ref_id": refID,
|
||
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued}},
|
||
}
|
||
now := domain.NowNano()
|
||
// list then update each (mon may not expose UpdateMany consistently)
|
||
var list []*domain.Job
|
||
if err := s.jobs.Find(ctx, &list, filter); err != nil {
|
||
return err
|
||
}
|
||
// Update is optimistic-concurrency guarded, so a lost race means the job is still pending and
|
||
// will run alongside its replacement. Swallowing that error reported a clean cancel that never
|
||
// happened; a stale job racing the new one is worth failing the call over.
|
||
for _, j := range list {
|
||
j.Status = domain.StatusCancelled
|
||
j.ProgressSummary = "已由新的定期排程取代"
|
||
j.CompletedAt = now
|
||
j.UpdatedAt = now
|
||
if err := s.Update(ctx, j); err != nil {
|
||
// 這一輪已經被別人搶去改（例如 worker 剛領走），不是我們該取消的目標。
|
||
if errors.Is(err, domain.ErrIllegalStatus) {
|
||
logx.WithContext(ctx).Infof("cancel pending job skipped, changed concurrently: job=%s", j.ID)
|
||
continue
|
||
}
|
||
return err
|
||
}
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *MonStore) CountActiveByOwnerAndTemplate(ctx context.Context, ownerUID int64, template string) (int64, error) {
|
||
filter := bson.M{
|
||
"owner_uid": ownerUID,
|
||
"template_type": template,
|
||
"status": bson.M{"$in": []string{domain.StatusPending, domain.StatusQueued, domain.StatusRunning}},
|
||
}
|
||
return s.jobs.CountDocuments(ctx, filter)
|
||
}
|
||
|
||
func (s *MonStore) Delete(ctx context.Context, id string) error {
|
||
res, err := s.jobs.DeleteOne(ctx, bson.M{"_id": id})
|
||
if err != nil {
|
||
return err
|
||
}
|
||
if res == 0 {
|
||
return domain.ErrNotFound
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (s *MonStore) DeleteTerminalBefore(ctx context.Context, beforeNs int64) (int64, error) {
|
||
// completed_at 優先；缺則用 updated_at（舊資料）
|
||
filter := bson.M{
|
||
"status": bson.M{"$in": []string{
|
||
domain.StatusSucceeded, domain.StatusFailed, domain.StatusCancelled,
|
||
}},
|
||
"$or": []bson.M{
|
||
{"completed_at": bson.M{"$gt": 0, "$lt": beforeNs}},
|
||
{
|
||
"$and": []bson.M{
|
||
{"$or": []bson.M{
|
||
{"completed_at": bson.M{"$exists": false}},
|
||
{"completed_at": 0},
|
||
}},
|
||
{"updated_at": bson.M{"$lt": beforeNs}},
|
||
},
|
||
},
|
||
},
|
||
}
|
||
// go-zero mon.DeleteMany returns deleted count (int64)
|
||
return s.jobs.DeleteMany(ctx, filter)
|
||
}
|
||
|
||
var _ domain.Repository = (*MonStore)(nil)
|