116 lines
3.4 KiB
Go
116 lines
3.4 KiB
Go
package job
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"hyapp-admin-server/internal/cache"
|
|
"hyapp-admin-server/internal/config"
|
|
"hyapp-admin-server/internal/model"
|
|
"hyapp-admin-server/internal/repository"
|
|
)
|
|
|
|
func TestArtifactJobRenewsDatabaseAndDistributedLeases(t *testing.T) {
|
|
store := &heartbeatStore{}
|
|
locker := &heartbeatLocker{}
|
|
runner := NewRunner(store, locker, config.Config{
|
|
NodeID: "worker-heartbeat",
|
|
Jobs: config.JobsConfig{LeaseTTL: 150 * time.Millisecond},
|
|
}, WithArtifactHandler("slow-export", func(ctx context.Context, _ *model.AdminJob) (string, string, error) {
|
|
select {
|
|
case <-ctx.Done():
|
|
return "", "", ctx.Err()
|
|
case <-time.After(350 * time.Millisecond):
|
|
return `{"rows":1}`, "slow.csv", nil
|
|
}
|
|
}))
|
|
|
|
if err := runner.handleWithHeartbeat(context.Background(), &model.AdminJob{ID: 91, Type: "slow-export"}); err != nil {
|
|
t.Fatalf("slow artifact handler failed: %v", err)
|
|
}
|
|
store.mu.Lock()
|
|
renews, completes, completedBy := store.renews, store.completes, store.completedBy
|
|
store.mu.Unlock()
|
|
locker.mu.Lock()
|
|
refreshes := locker.refreshes
|
|
locker.mu.Unlock()
|
|
if renews < 2 || refreshes < 2 || completes != 1 || completedBy != "worker-heartbeat" {
|
|
t.Fatalf("long job lease was not kept alive: renews=%d refreshes=%d completes=%d worker=%q", renews, refreshes, completes, completedBy)
|
|
}
|
|
}
|
|
|
|
func TestArtifactJobCancelsWhenDistributedLeaseIsLost(t *testing.T) {
|
|
store := &heartbeatStore{}
|
|
locker := &heartbeatLocker{loseOnRefresh: true}
|
|
runner := NewRunner(store, locker, config.Config{
|
|
NodeID: "worker-lost-lock",
|
|
Jobs: config.JobsConfig{LeaseTTL: 150 * time.Millisecond},
|
|
}, WithArtifactHandler("blocked-export", func(ctx context.Context, _ *model.AdminJob) (string, string, error) {
|
|
<-ctx.Done()
|
|
return "", "", ctx.Err()
|
|
}))
|
|
|
|
err := runner.handleWithHeartbeat(context.Background(), &model.AdminJob{ID: 92, Type: "blocked-export"})
|
|
if !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("lost distributed lock must cancel handler, err=%v", err)
|
|
}
|
|
store.mu.Lock()
|
|
completes := store.completes
|
|
store.mu.Unlock()
|
|
if completes != 0 {
|
|
t.Fatalf("worker that lost its lease must not complete the job: completes=%d", completes)
|
|
}
|
|
}
|
|
|
|
type heartbeatStore struct {
|
|
mu sync.Mutex
|
|
renews int
|
|
completes int
|
|
completedBy string
|
|
}
|
|
|
|
func (s *heartbeatStore) LeaseNextJob(string, time.Duration, int) (*model.AdminJob, error) {
|
|
return nil, errors.New("not used")
|
|
}
|
|
|
|
func (s *heartbeatStore) RenewJobLease(_ uint, _ string, _ time.Duration) error {
|
|
s.mu.Lock()
|
|
s.renews++
|
|
s.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (s *heartbeatStore) CompleteJobWithArtifact(_ uint, workerID string, _, _ string) error {
|
|
s.mu.Lock()
|
|
s.completes++
|
|
s.completedBy = workerID
|
|
s.mu.Unlock()
|
|
return nil
|
|
}
|
|
|
|
func (s *heartbeatStore) FailJobAttempt(uint, string, string) error { return nil }
|
|
func (s *heartbeatStore) ReleaseJobLease(uint, string) error { return nil }
|
|
func (s *heartbeatStore) ExportUsers(repository.ListOptions) ([]model.User, error) {
|
|
return nil, nil
|
|
}
|
|
|
|
type heartbeatLocker struct {
|
|
mu sync.Mutex
|
|
refreshes int
|
|
loseOnRefresh bool
|
|
}
|
|
|
|
func (l *heartbeatLocker) TryLock(context.Context, string, string, time.Duration) (bool, cache.UnlockFunc, error) {
|
|
return true, func(context.Context) error { return nil }, nil
|
|
}
|
|
|
|
func (l *heartbeatLocker) RefreshLock(context.Context, string, string, time.Duration) (bool, error) {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
l.refreshes++
|
|
return !l.loseOnRefresh, nil
|
|
}
|