projecteru2 · CMGS · Jun 27, 2023 · Jun 21, 2023 · Jun 26, 2023
diff --git a/cluster/calcium/control.go b/cluster/calcium/control.go
@@ -43,6 +43,12 @@ func (c *Calcium) ControlWorkload(ctx context.Context, IDs []string, typ string,
 						startHook, err := c.doStartWorkload(ctx, workload, force)
 						message = append(message, startHook...)
 						return err
+					case cluster.WorkloadSuspend:
+						message, err = c.doSuspendWorkload(ctx, workload, force)
+						return err
+					case cluster.WorkloadResume:
+						message, err = c.doResumeWorkload(ctx, workload, force)
+						return err
 					}
 					return types.ErrInvaildControlType
 				})
@@ -103,3 +109,39 @@ func (c *Calcium) doStopWorkload(ctx context.Context, workload *types.Workload,
 	}
 	return message, err
 }
+
+func (c *Calcium) doSuspendWorkload(ctx context.Context, workload *types.Workload, force bool) (message []*bytes.Buffer, err error) {
+	if workload.Hook != nil && len(workload.Hook.BeforeSuspend) > 0 {
+		message, err = c.doHook(
+			ctx,
+			workload.ID, workload.User,
+			workload.Hook.BeforeSuspend, workload.Env,
+			workload.Hook.Force, workload.Privileged,
+			force, workload.Engine,
+		)
+		if err != nil {
+			return message, err
+		}
+	}
+
+	if err = workload.Suspend(ctx); err != nil {
+		message = append(message, bytes.NewBufferString(err.Error()))
+	}
+	return message, err
+}
+
+func (c *Calcium) doResumeWorkload(ctx context.Context, workload *types.Workload, force bool) (message []*bytes.Buffer, err error) {
+	if err = workload.Resume(ctx); err != nil {
+		return message, err
+	}
+	if workload.Hook != nil && len(workload.Hook.AfterResume) > 0 {
+		message, err = c.doHook(
+			ctx,
+			workload.ID, workload.User,
+			workload.Hook.AfterResume, workload.Env,
+			workload.Hook.Force, workload.Privileged,
+			force, workload.Engine,
+		)
+	}
+	return message, err
+}
diff --git a/cluster/calcium/control_test.go b/cluster/calcium/control_test.go
@@ -178,3 +178,128 @@ func TestControlRestart(t *testing.T) {
 		assert.NoError(t, r.Error)
 	}
 }
+
+func TestControlSuspend(t *testing.T) {
+	c := NewTestCluster()
+	ctx := context.Background()
+	store := c.store.(*storemocks.Store)
+	lock := &lockmocks.DistributedLock{}
+	lock.On("Lock", mock.Anything).Return(ctx, nil)
+	lock.On("Unlock", mock.Anything).Return(nil)
+	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
+	workload := &types.Workload{
+		ID:         "id1",
+		Privileged: true,
+	}
+	engine := &enginemocks.API{}
+	workload.Engine = engine
+	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
+	// failed, hook true, remove always false
+	hook := &types.Hook{
+		BeforeSuspend: []string{"cmd1"},
+	}
+	workload.Hook = hook
+	workload.Hook.Force = true
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("", nil, nil, nil, types.ErrNilEngine)
+	ch, err := c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadSuspend, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	// stop failed
+	workload.Hook.Force = false
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadSuspend, false)
+	engine.On("VirtualizationSuspend", mock.Anything, mock.Anything, mock.Anything).Return(types.ErrNilEngine).Once()
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	engine.On("VirtualizationSuspend", mock.Anything, mock.Anything, mock.Anything).Return(nil)
+	// stop success
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadSuspend, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.NoError(t, r.Error)
+	}
+}
+
+func TestControlResume(t *testing.T) {
+	c := NewTestCluster()
+	ctx := context.Background()
+	store := c.store.(*storemocks.Store)
+	lock := &lockmocks.DistributedLock{}
+	lock.On("Lock", mock.Anything).Return(ctx, nil)
+	lock.On("Unlock", mock.Anything).Return(nil)
+	store.On("CreateLock", mock.Anything, mock.Anything).Return(lock, nil)
+	// failed by GetWorkloads
+	store.On("GetWorkloads", mock.Anything, mock.Anything).Return(nil, types.ErrMockError).Once()
+	ch, err := c.ControlWorkload(ctx, []string{"id1"}, "", true)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	workload := &types.Workload{
+		ID:         "id1",
+		Privileged: true,
+	}
+	engine := &enginemocks.API{}
+	workload.Engine = engine
+	store.On("GetWorkloads", mock.Anything, mock.Anything).Return([]*types.Workload{workload}, nil)
+	// failed by type
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, "", true)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	// failed by start
+	engine.On("VirtualizationResume", mock.Anything, mock.Anything).Return(types.ErrNilEngine).Once()
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadResume, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	engine.On("VirtualizationResume", mock.Anything, mock.Anything).Return(nil)
+	// failed by Execute
+	hook := &types.Hook{
+		AfterResume: []string{"cmd1", "cmd2"},
+	}
+	workload.Hook = hook
+	workload.Hook.Force = false
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("", nil, nil, nil, types.ErrNilEngine).Times(3)
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadResume, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.NoError(t, r.Error)
+	}
+	// force false, get no error
+	workload.Hook.Force = true
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadResume, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+		assert.Equal(t, r.WorkloadID, "id1")
+	}
+	data := io.NopCloser(bytes.NewBufferString("output"))
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("eid", data, nil, nil, nil).Times(4)
+	// failed by ExecExitCode
+	engine.On("ExecExitCode", mock.Anything, mock.Anything, mock.Anything).Return(-1, types.ErrNilEngine).Once()
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadResume, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	// exitCode is not 0
+	engine.On("ExecExitCode", mock.Anything, mock.Anything, mock.Anything).Return(-1, nil).Once()
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadResume, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.Error(t, r.Error)
+	}
+	// exitCode is 0
+	engine.On("ExecExitCode", mock.Anything, mock.Anything, mock.Anything).Return(0, nil)
+	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadResume, false)
+	assert.NoError(t, err)
+	for r := range ch {
+		assert.NoError(t, r.Error)
+	}
+}
diff --git a/cluster/cluster.go b/cluster/cluster.go
@@ -28,6 +28,10 @@ const (
 	WorkloadStart = "start"
 	// WorkloadRestart for restart workload
 	WorkloadRestart = "restart"
+	// WorkloadSuspend for suspending workload
+	WorkloadSuspend = "suspend"
+	// WorkloadResume for resuming workload
+	WorkloadResume = "resume"
 	// WorkloadLock for lock workload
 	WorkloadLock = "clock_%s"
 	// PodLock for lock pod

diff --git a/engine/docker/container.go b/engine/docker/container.go
@@ -307,6 +307,16 @@ func (e *Engine) VirtualizationStop(ctx context.Context, ID string, gracefulTime
 	return e.client.ContainerStop(ctx, ID, dockercontainer.StopOptions{Timeout: timeout})
 }
 
+// VirtualizationSuspend suspends virtualization
+func (e *Engine) VirtualizationSuspend(context.Context, string) error {
+	return nil
+}
+
+// VirtualizationResume resumes virtualization
+func (e *Engine) VirtualizationResume(context.Context, string) error {
+	return nil
+}
+
 // VirtualizationRemove remove virtualization
 func (e *Engine) VirtualizationRemove(ctx context.Context, ID string, removeVolumes, force bool) error {
 	if err := e.client.ContainerRemove(ctx, ID, dockertypes.ContainerRemoveOptions{RemoveVolumes: removeVolumes, Force: force}); err != nil {

diff --git a/engine/engine.go b/engine/engine.go
@@ -43,6 +43,8 @@ type API interface {
 	VirtualizationStart(ctx context.Context, ID string) error
 	VirtualizationStop(ctx context.Context, ID string, gracefulTimeout time.Duration) error
 	VirtualizationRemove(ctx context.Context, ID string, volumes, force bool) error
+	VirtualizationSuspend(ctx context.Context, ID string) error
+	VirtualizationResume(ctx context.Context, ID string) error
 	VirtualizationInspect(ctx context.Context, ID string) (*enginetypes.VirtualizationInfo, error)
 	VirtualizationLogs(ctx context.Context, opts *enginetypes.VirtualizationLogStreamOptions) (stdout, stderr io.ReadCloser, err error)
 	VirtualizationAttach(ctx context.Context, ID string, stream, openStdin bool) (stdout, stderr io.ReadCloser, stdin io.WriteCloser, err error)

diff --git a/engine/fake/fake.go b/engine/fake/fake.go
@@ -140,6 +140,16 @@ func (f *EngineWithErr) VirtualizationStop(context.Context, string, time.Duratio
 	return f.DefaultErr
 }
 
+// VirtualizationSuspend .
+func (f *EngineWithErr) VirtualizationSuspend(context.Context, string) error {
+	return f.DefaultErr
+}
+
+// VirtualizationResume .
+func (f *EngineWithErr) VirtualizationResume(context.Context, string) error {
+	return f.DefaultErr
+}
+
 // VirtualizationRemove .
 func (f *EngineWithErr) VirtualizationRemove(context.Context, string, bool, bool) error {
 	return f.DefaultErr