github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/cluster/calcium/control.go (about)

     1  package calcium
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"sync"
     7  
     8  	"github.com/projecteru2/core/cluster"
     9  	"github.com/projecteru2/core/log"
    10  	"github.com/projecteru2/core/types"
    11  	"github.com/projecteru2/core/utils"
    12  )
    13  
    14  // ControlWorkload control workloads status
    15  func (c *Calcium) ControlWorkload(ctx context.Context, IDs []string, typ string, force bool) (chan *types.ControlWorkloadMessage, error) {
    16  	logger := log.WithFunc("calcium.ControlWorkload").WithField("IDs", IDs).WithField("typ", typ).WithField("force", force)
    17  	ch := make(chan *types.ControlWorkloadMessage)
    18  
    19  	_ = c.pool.Invoke(func() {
    20  		defer close(ch)
    21  		wg := &sync.WaitGroup{}
    22  		wg.Add(len(IDs))
    23  		defer wg.Wait()
    24  		for _, ID := range IDs {
    25  			ID := ID
    26  			_ = c.pool.Invoke(func() {
    27  				defer wg.Done()
    28  				var message []*bytes.Buffer
    29  				err := c.withWorkloadLocked(ctx, ID, false, func(ctx context.Context, workload *types.Workload) error {
    30  					var err error
    31  					switch typ {
    32  					case cluster.WorkloadStop:
    33  						message, err = c.doStopWorkload(ctx, workload, force)
    34  						return err
    35  					case cluster.WorkloadStart:
    36  						message, err = c.doStartWorkload(ctx, workload, force)
    37  						return err
    38  					case cluster.WorkloadRestart:
    39  						message, err = c.doStopWorkload(ctx, workload, force)
    40  						if err != nil {
    41  							return err
    42  						}
    43  						startHook, err := c.doStartWorkload(ctx, workload, force)
    44  						message = append(message, startHook...)
    45  						return err
    46  					case cluster.WorkloadSuspend:
    47  						message, err = c.doSuspendWorkload(ctx, workload, force)
    48  						return err
    49  					case cluster.WorkloadResume:
    50  						message, err = c.doResumeWorkload(ctx, workload, force)
    51  						return err
    52  					}
    53  					return types.ErrInvaildControlType
    54  				})
    55  				if err == nil {
    56  					logger.Infof(ctx, "Workload %s %s", ID, typ)
    57  					logger.Infof(ctx, "%+v", "Hook Output:")
    58  					logger.Infof(ctx, "%+v", string(utils.MergeHookOutputs(message)))
    59  				}
    60  				logger.Error(ctx, err)
    61  				ch <- &types.ControlWorkloadMessage{
    62  					WorkloadID: ID,
    63  					Error:      err,
    64  					Hook:       message,
    65  				}
    66  			})
    67  		}
    68  	})
    69  
    70  	return ch, nil
    71  }
    72  
    73  func (c *Calcium) doStartWorkload(ctx context.Context, workload *types.Workload, force bool) (message []*bytes.Buffer, err error) {
    74  	if err = workload.Start(ctx); err != nil {
    75  		return message, err
    76  	}
    77  	// TODO healthcheck first
    78  	if workload.Hook != nil && len(workload.Hook.AfterStart) > 0 {
    79  		message, err = c.doHook(
    80  			ctx,
    81  			workload.ID, workload.User,
    82  			workload.Hook.AfterStart, workload.Env,
    83  			workload.Hook.Force, workload.Privileged,
    84  			force, workload.Engine,
    85  		)
    86  	}
    87  	return message, err
    88  }
    89  
    90  func (c *Calcium) doStopWorkload(ctx context.Context, workload *types.Workload, force bool) (message []*bytes.Buffer, err error) {
    91  	if workload.Hook != nil && len(workload.Hook.BeforeStop) > 0 {
    92  		message, err = c.doHook(
    93  			ctx,
    94  			workload.ID, workload.User,
    95  			workload.Hook.BeforeStop, workload.Env,
    96  			workload.Hook.Force, workload.Privileged,
    97  			force, workload.Engine,
    98  		)
    99  		if err != nil {
   100  			return message, err
   101  		}
   102  	}
   103  
   104  	// 这里 block 的问题很严重，按照目前的配置是 5 分钟一级的 block
   105  	// 一个简单的处理方法是相信 ctx 不相信 engine 自身的处理
   106  	// 另外我怀疑 engine 自己的 timeout 实现是完全的等 timeout 而非结束了就退出
   107  	if err = workload.Stop(ctx, force); err != nil {
   108  		message = append(message, bytes.NewBufferString(err.Error()))
   109  	}
   110  	return message, err
   111  }
   112  
   113  func (c *Calcium) doSuspendWorkload(ctx context.Context, workload *types.Workload, force bool) (message []*bytes.Buffer, err error) {
   114  	if workload.Hook != nil && len(workload.Hook.BeforeSuspend) > 0 {
   115  		message, err = c.doHook(
   116  			ctx,
   117  			workload.ID, workload.User,
   118  			workload.Hook.BeforeSuspend, workload.Env,
   119  			workload.Hook.Force, workload.Privileged,
   120  			force, workload.Engine,
   121  		)
   122  		if err != nil {
   123  			return message, err
   124  		}
   125  	}
   126  
   127  	if err = workload.Suspend(ctx); err != nil {
   128  		message = append(message, bytes.NewBufferString(err.Error()))
   129  	}
   130  	return message, err
   131  }
   132  
   133  func (c *Calcium) doResumeWorkload(ctx context.Context, workload *types.Workload, force bool) (message []*bytes.Buffer, err error) {
   134  	if err = workload.Resume(ctx); err != nil {
   135  		return message, err
   136  	}
   137  	if workload.Hook != nil && len(workload.Hook.AfterResume) > 0 {
   138  		message, err = c.doHook(
   139  			ctx,
   140  			workload.ID, workload.User,
   141  			workload.Hook.AfterResume, workload.Env,
   142  			workload.Hook.Force, workload.Privileged,
   143  			force, workload.Engine,
   144  		)
   145  	}
   146  	return message, err
   147  }