Skip to content

Commit

Permalink
feat(restore): adds exclusive cluster lock for restore task (#4122)
Browse files Browse the repository at this point in the history 
* feat(restore): adds exclusive cluster lock for restore task

This makes the scheduler run the restore task only if no other tasks (backup
or repair) are running in the cluster. Conversely, other tasks (backup or
repair) won't start if the restore task is running in the cluster at that
time.

Fixes: #4045.
  • Loading branch information
@VAveryanov8
VAveryanov8 authored Nov 26, 2024
1 parent 75fb75c commit b97d1cf
Show file tree
Hide file tree
Showing 4 changed files with 193 additions and 58 deletions.
8 changes: 5 additions & 3 deletions pkg/cmd/scylla-manager/server.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -144,11 +144,13 @@ func (s *server) makeServices(ctx context.Context) error {
return errors.Wrapf(err, "scheduler service")
}

restoreExclusiveLock := scheduler.NewTaskExclusiveLockPolicy(scheduler.RestoreTask)

// Register the runners
s.schedSvc.SetRunner(scheduler.BackupTask, scheduler.PolicyRunner{Policy: scheduler.NewLockClusterPolicy(), Runner: s.backupSvc.Runner()})
s.schedSvc.SetRunner(scheduler.RestoreTask, scheduler.PolicyRunner{Policy: scheduler.NewLockClusterPolicy(), Runner: s.restoreSvc.Runner()})
s.schedSvc.SetRunner(scheduler.BackupTask, scheduler.PolicyRunner{Policy: restoreExclusiveLock, Runner: s.backupSvc.Runner(), TaskType: scheduler.BackupTask})
s.schedSvc.SetRunner(scheduler.RestoreTask, scheduler.PolicyRunner{Policy: restoreExclusiveLock, Runner: s.restoreSvc.Runner(), TaskType: scheduler.RestoreTask})
s.schedSvc.SetRunner(scheduler.HealthCheckTask, s.healthSvc.Runner())
s.schedSvc.SetRunner(scheduler.RepairTask, scheduler.PolicyRunner{Policy: scheduler.NewLockClusterPolicy(), Runner: s.repairSvc.Runner()})
s.schedSvc.SetRunner(scheduler.RepairTask, scheduler.PolicyRunner{Policy: restoreExclusiveLock, Runner: s.repairSvc.Runner(), TaskType: scheduler.RepairTask})
s.schedSvc.SetRunner(scheduler.ValidateBackupTask, s.backupSvc.ValidationRunner())

// Add additional properties on task run.
Expand Down
16 changes: 8 additions & 8 deletions pkg/service/scheduler/mock_policy_test.go
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

92 changes: 70 additions & 22 deletions pkg/service/scheduler/policy.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ package scheduler
import (
"context"
"encoding/json"
"fmt"
"sync"

"github.com/pkg/errors"
Expand All @@ -13,55 +14,102 @@ import (

// Policy decides if given task can be run.
type Policy interface {
PreRun(clusterID, taskID, runID uuid.UUID) error
PostRun(clusterID, taskID, runID uuid.UUID)
PreRun(clusterID, taskID, runID uuid.UUID, taskType TaskType) error
PostRun(clusterID, taskID, runID uuid.UUID, taskType TaskType)
}

// PolicyRunner is a runner that uses policy to check if a task can be run.
type PolicyRunner struct {
Policy Policy
Runner Runner

// TaskType of a task that this runner is executing
TaskType TaskType
}

// Run implements Runner.
func (pr PolicyRunner) Run(ctx context.Context, clusterID, taskID, runID uuid.UUID, properties json.RawMessage) error {
if err := pr.Policy.PreRun(clusterID, taskID, runID); err != nil {
if err := pr.Policy.PreRun(clusterID, taskID, runID, pr.TaskType); err != nil {
return err
}
defer pr.Policy.PostRun(clusterID, taskID, runID)
defer pr.Policy.PostRun(clusterID, taskID, runID, pr.TaskType)
return pr.Runner.Run(ctx, clusterID, taskID, runID, properties)
}

var errClusterBusy = errors.New("another task is running")

// LockClusterPolicy is a policy that can execute only one task at a time.
type LockClusterPolicy struct {
mu sync.Mutex
busy map[uuid.UUID]struct{}
// TaskExclusiveLockPolicy is a policy that executes the exclusiveTask only if there are no other tasks in the cluster.
// Conversely, other tasks can run only if the exclusiveTask is not running.
// Additionally this policy ensures that only one task of a task type can be executed at a time in a cluster.
type TaskExclusiveLockPolicy struct {
mu sync.Mutex
running map[uuid.UUID]map[TaskType]struct{}

exclusiveTask TaskType
}

func NewLockClusterPolicy() *LockClusterPolicy {
return &LockClusterPolicy{
busy: make(map[uuid.UUID]struct{}),
func NewTaskExclusiveLockPolicy(exclusiveTask TaskType) *TaskExclusiveLockPolicy {
return &TaskExclusiveLockPolicy{
running: map[uuid.UUID]map[TaskType]struct{}{},

exclusiveTask: exclusiveTask,
}
}

// PreRun implements Policy.
func (p *LockClusterPolicy) PreRun(clusterID, _, _ uuid.UUID) error {
p.mu.Lock()
defer p.mu.Unlock()
// PreRun acquires exclusive lock on a cluster for a provided taskType.
func (t *TaskExclusiveLockPolicy) PreRun(clusterID, _, _ uuid.UUID, taskType TaskType) error {
t.mu.Lock()
defer t.mu.Unlock()

cluster, ok := t.running[clusterID]
if !ok {
cluster = map[TaskType]struct{}{}
t.running[clusterID] = cluster
}

if err := t.canRunTaskExclusively(cluster, taskType); err != nil {
// cluster is busy
return err
}

if _, ok := p.busy[clusterID]; ok {
cluster[taskType] = struct{}{}

return nil
}

// canRunTaskExclusively returns nil if taskType can be run in the cluster, otherwise err is returned.
func (t *TaskExclusiveLockPolicy) canRunTaskExclusively(cluster map[TaskType]struct{}, taskType TaskType) error {
// No tasks are running, so we can start the task.
if len(cluster) == 0 {
return nil
}

// Exclusive task can be run only when no other tasks is running.
if taskType == t.exclusiveTask {
return fmt.Errorf("run exclusive task %s: %w", taskType, errClusterBusy)
}

// Any other task can't be run when exclusive task is running.
if _, ok := cluster[t.exclusiveTask]; ok {
return fmt.Errorf("exclusive task (%s) is running: %w", taskType, errClusterBusy)
}

// Only one task of a taskType can run in a cluster at a time.
if _, ok := cluster[taskType]; ok {
return errClusterBusy
}

p.busy[clusterID] = struct{}{}
return nil
}

// PostRun implements Policy.
func (p *LockClusterPolicy) PostRun(clusterID, _, _ uuid.UUID) {
p.mu.Lock()
defer p.mu.Unlock()
delete(p.busy, clusterID)
// PostRun releases a lock on a cluster for a provided taskType.
func (t *TaskExclusiveLockPolicy) PostRun(clusterID, _, _ uuid.UUID, taskType TaskType) {
t.mu.Lock()
defer t.mu.Unlock()
// No need to check if t.running[clusterID] exists, because built-in delete will no-op in case of nil map.
delete(t.running[clusterID], taskType)
// Cleaning up the map if no more tasks left in the cluster.
if len(t.running[clusterID]) == 0 {
delete(t.running, clusterID)
}
}
135 changes: 110 additions & 25 deletions pkg/service/scheduler/policy_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,15 +22,17 @@ func TestPolicyRunner(t *testing.T) {
c := uuid.MustRandom()
k := uuid.MustRandom()
r := uuid.MustRandom()
tt := TaskType("")
e := errors.New("test")

mp := NewmockPolicy(ctrl)
mp.EXPECT().PreRun(c, k, r).Return(e)
mp.EXPECT().PreRun(c, k, r, tt).Return(e)
mr := NewmockRunner(ctrl)

p := PolicyRunner{
Policy: mp,
Runner: mr,
Policy: mp,
Runner: mr,
TaskType: tt,
}
if err := p.Run(context.Background(), c, k, r, nil); err != e {
t.Fatal("expected", e, "got", err)
Expand All @@ -44,47 +46,130 @@ func TestPolicyRunner(t *testing.T) {
c := uuid.MustRandom()
k := uuid.MustRandom()
r := uuid.MustRandom()
tt := TaskType("")
e := errors.New("test")

mp := NewmockPolicy(ctrl)
gomock.InOrder(
mp.EXPECT().PreRun(c, k, r).Return(nil),
mp.EXPECT().PostRun(c, k, r),
mp.EXPECT().PreRun(c, k, r, tt).Return(nil),
mp.EXPECT().PostRun(c, k, r, tt),
)
mr := NewmockRunner(ctrl)
mr.EXPECT().Run(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(e)

p := PolicyRunner{
Policy: mp,
Runner: mr,
Policy: mp,
Runner: mr,
TaskType: tt,
}
if err := p.Run(context.Background(), c, k, r, nil); err != e {
t.Fatal("expected", e, "got", err)
}
})
}

func TestNewLockClusterPolicy(t *testing.T) {
c := uuid.MustRandom()
k := uuid.MustRandom()
r := uuid.MustRandom()
p := NewLockClusterPolicy()
func TestExclusiveTaskLockPolicy(t *testing.T) {
clusterID := uuid.MustRandom()
runID := uuid.MustRandom()
taskID := uuid.MustRandom()

if err := p.PreRun(c, k, r); err != nil {
t.Fatal(err)
}
t.Run("when no other task is running, preRun should return nil", func(t *testing.T) {
restoreExclusiveTask := NewTaskExclusiveLockPolicy(RestoreTask)

if err := p.PreRun(c, uuid.MustRandom(), uuid.MustRandom()); err == nil {
t.Fatal("expected error")
}
err := restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}
})

t.Run("when exclusive task is running, other tasks are not allowed", func(t *testing.T) {
restoreExclusiveTask := NewTaskExclusiveLockPolicy(RestoreTask)

err := restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}

err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, BackupTask)
if !errors.Is(err, errClusterBusy) {
t.Fatalf("PreRun: expected errClusterBusy, got: %v", err)
}
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, RepairTask)
if !errors.Is(err, errClusterBusy) {
t.Fatalf("PreRun: expected errClusterBusy, got: %v", err)
}
})

t.Run("when non exclusive task is running, exclusive task is not allowed", func(t *testing.T) {
restoreExclusiveTask := NewTaskExclusiveLockPolicy(RestoreTask)

err := restoreExclusiveTask.PreRun(clusterID, taskID, runID, BackupTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, RepairTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}

err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if !errors.Is(err, errClusterBusy) {
t.Fatalf("PreRun: expected errClusterBusy, got: %v", err)
}
})

t.Run("only one instance of a task type is allowed to run at a time", func(t *testing.T) {
restoreExclusiveTask := NewTaskExclusiveLockPolicy(RestoreTask)

if err := p.PreRun(uuid.MustRandom(), uuid.MustRandom(), uuid.MustRandom()); err != nil {
t.Fatal(err)
}
err := restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if !errors.Is(err, errClusterBusy) {
t.Fatalf("PreRun: expected errClusterBusy, got: %v", err)
}

restoreExclusiveTask = NewTaskExclusiveLockPolicy(RestoreTask)
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, BackupTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, BackupTask)
if !errors.Is(err, errClusterBusy) {
t.Fatalf("PreRun: expected errClusterBusy, got: %v", err)
}
})

t.Run("PostRun on a empty cluster", func(t *testing.T) {
restoreExclusiveTask := NewTaskExclusiveLockPolicy(RestoreTask)

restoreExclusiveTask.PostRun(clusterID, taskID, runID, RestoreTask)
})

p.PostRun(c, uuid.MustRandom(), uuid.MustRandom())
t.Run("PostRun should release lock for a given task type", func(t *testing.T) {
restoreExclusiveTask := NewTaskExclusiveLockPolicy(RestoreTask)

if err := p.PreRun(c, uuid.MustRandom(), uuid.MustRandom()); err != nil {
t.Fatal(errClusterBusy)
}
err := restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if !errors.Is(err, errClusterBusy) {
t.Fatalf("PreRun: expected errClusterBusy, got: %v", err)
}

// Release a lock and clean up the underlying map.
restoreExclusiveTask.PostRun(clusterID, taskID, runID, RestoreTask)

if _, ok := restoreExclusiveTask.running[clusterID]; ok {
t.Fatalf("t.running[clusterID] should be deleted")
}

// Lock can be acquried again.
err = restoreExclusiveTask.PreRun(clusterID, taskID, runID, RestoreTask)
if err != nil {
t.Fatalf("PreRun: unexpected err: %v", err)
}
})
}

0 comments on commit b97d1cf

Please sign in to comment.