github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/taskrunner/state/state.go (about)

     1  package state
     2  
     3  import (
     4  	"github.com/hashicorp/nomad/plugins/drivers"
     5  	"golang.org/x/exp/maps"
     6  )
     7  
     8  // LocalState is Task state which is persisted for use when restarting Nomad
     9  // agents.
    10  type LocalState struct {
    11  	Hooks map[string]*HookState
    12  
    13  	// DriverNetwork is the network information returned by the task
    14  	// driver's Start method
    15  	DriverNetwork *drivers.DriverNetwork
    16  
    17  	// TaskHandle is the handle used to reattach to the task during recovery
    18  	TaskHandle *drivers.TaskHandle
    19  
    20  	// RunComplete is set to true when the TaskRunner.Run() method finishes.
    21  	// It is used to distinguish between a dead task that could be restarted
    22  	// and one that will never run again.
    23  	RunComplete bool
    24  }
    25  
    26  func NewLocalState() *LocalState {
    27  	return &LocalState{
    28  		Hooks: make(map[string]*HookState),
    29  	}
    30  }
    31  
    32  // Canonicalize ensures LocalState is in a consistent state by initializing
    33  // Hooks and ensuring no HookState's are nil. Useful for cleaning unmarshalled
    34  // state which may be in an unknown state.
    35  func (s *LocalState) Canonicalize() {
    36  	if s.Hooks == nil {
    37  		// Hooks is nil, create it
    38  		s.Hooks = make(map[string]*HookState)
    39  	} else {
    40  		for k, v := range s.Hooks {
    41  			// Remove invalid nil entries from Hooks map
    42  			if v == nil {
    43  				delete(s.Hooks, k)
    44  			}
    45  		}
    46  	}
    47  }
    48  
    49  // Copy LocalState. Returns nil if nil.
    50  func (s *LocalState) Copy() *LocalState {
    51  	if s == nil {
    52  		return nil
    53  	}
    54  
    55  	// Create a copy
    56  	c := &LocalState{
    57  		Hooks:         make(map[string]*HookState, len(s.Hooks)),
    58  		DriverNetwork: s.DriverNetwork.Copy(),
    59  		TaskHandle:    s.TaskHandle.Copy(),
    60  		RunComplete:   s.RunComplete,
    61  	}
    62  
    63  	// Copy the hook state
    64  	for h, state := range s.Hooks {
    65  		c.Hooks[h] = state.Copy()
    66  	}
    67  
    68  	return c
    69  }
    70  
    71  type HookState struct {
    72  	// Prestart is true if the hook has run Prestart successfully and does
    73  	// not need to run again
    74  	PrestartDone bool
    75  
    76  	// Data allows hooks to persist arbitrary state.
    77  	Data map[string]string
    78  
    79  	// Environment variables set by the hook that will continue to be set
    80  	// even if PrestartDone=true.
    81  	Env map[string]string
    82  }
    83  
    84  // Copy HookState. Returns nil if its nil.
    85  func (h *HookState) Copy() *HookState {
    86  	if h == nil {
    87  		return nil
    88  	}
    89  
    90  	c := new(HookState)
    91  	*c = *h
    92  	c.Data = maps.Clone(h.Data)
    93  	c.Env = maps.Clone(h.Env)
    94  	return c
    95  }
    96  
    97  func (h *HookState) Equal(o *HookState) bool {
    98  	if h == nil || o == nil {
    99  		return h == o
   100  	}
   101  
   102  	if h.PrestartDone != o.PrestartDone {
   103  		return false
   104  	}
   105  
   106  	if !maps.Equal(h.Data, o.Data) {
   107  		return false
   108  	}
   109  
   110  	return maps.Equal(h.Env, o.Env)
   111  }