github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/worker/caasunitterminationworker/worker.go (about)

     1  // Copyright 2021 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package caasunitterminationworker
     5  
     6  import (
     7  	"os"
     8  	"os/signal"
     9  	"syscall"
    10  
    11  	"github.com/juju/clock"
    12  	"github.com/juju/errors"
    13  	"github.com/juju/names/v5"
    14  	"github.com/juju/worker/v3"
    15  	"gopkg.in/tomb.v2"
    16  
    17  	"github.com/juju/juju/agent"
    18  	"github.com/juju/juju/api/agent/caasapplication"
    19  )
    20  
    21  // TerminationSignal is SIGTERM which is sent by most container runtimes when
    22  // a container should terminate gracefully.
    23  const TerminationSignal = syscall.SIGTERM
    24  
    25  type terminationWorker struct {
    26  	tomb tomb.Tomb
    27  
    28  	agent          agent.Agent
    29  	state          State
    30  	unitTerminator UnitTerminator
    31  	logger         Logger
    32  	clock          clock.Clock
    33  }
    34  
    35  type Config struct {
    36  	Agent          agent.Agent
    37  	State          State
    38  	UnitTerminator UnitTerminator
    39  	Logger         Logger
    40  	Clock          clock.Clock
    41  }
    42  
    43  type State interface {
    44  	UnitTerminating(tag names.UnitTag) (caasapplication.UnitTermination, error)
    45  }
    46  
    47  type UnitTerminator interface {
    48  	Terminate() error
    49  }
    50  
    51  // NewWorker returns a worker that waits for a
    52  // TerminationSignal signal, and then exits
    53  // with worker.ErrTerminateAgent.
    54  func NewWorker(config Config) worker.Worker {
    55  	w := terminationWorker{
    56  		agent:          config.Agent,
    57  		state:          config.State,
    58  		unitTerminator: config.UnitTerminator,
    59  		logger:         config.Logger,
    60  		clock:          config.Clock,
    61  	}
    62  	c := make(chan os.Signal, 1)
    63  	signal.Notify(c, TerminationSignal)
    64  	w.tomb.Go(func() error {
    65  		defer signal.Stop(c)
    66  		return w.loop(c)
    67  	})
    68  	return &w
    69  }
    70  
    71  func (w *terminationWorker) Kill() {
    72  	w.tomb.Kill(nil)
    73  }
    74  
    75  func (w *terminationWorker) Wait() error {
    76  	return w.tomb.Wait()
    77  }
    78  
    79  func (w *terminationWorker) loop(c <-chan os.Signal) (err error) {
    80  	select {
    81  	case <-c:
    82  		w.logger.Infof("terminating due to SIGTERM")
    83  		term, err := w.state.UnitTerminating(w.agent.CurrentConfig().Tag().(names.UnitTag))
    84  		if err != nil {
    85  			w.logger.Errorf("error while terminating unit: %v", err)
    86  			return err
    87  		}
    88  		if !term.WillRestart {
    89  			// Lifecycle watcher will handle termination of the unit.
    90  			return nil
    91  		}
    92  		err = w.unitTerminator.Terminate()
    93  		if err != nil {
    94  			w.logger.Errorf("error while terminating unit: %v", err)
    95  			return errors.Annotatef(err, "failed to terminate unit agent worker")
    96  		}
    97  		return nil
    98  	case <-w.tomb.Dying():
    99  		return tomb.ErrDying
   100  	}
   101  }