github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocrunner/taskrunner/state/state.go (about)

     1  package state
     2  
     3  import (
     4  	"github.com/hashicorp/nomad/helper"
     5  	"github.com/hashicorp/nomad/plugins/drivers"
     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  
    21  func NewLocalState() *LocalState {
    22  	return &LocalState{
    23  		Hooks: make(map[string]*HookState),
    24  	}
    25  }
    26  
    27  // Canonicalize ensures LocalState is in a consistent state by initializing
    28  // Hooks and ensuring no HookState's are nil. Useful for cleaning unmarshalled
    29  // state which may be in an unknown state.
    30  func (s *LocalState) Canonicalize() {
    31  	if s.Hooks == nil {
    32  		// Hooks is nil, create it
    33  		s.Hooks = make(map[string]*HookState)
    34  	} else {
    35  		for k, v := range s.Hooks {
    36  			// Remove invalid nil entries from Hooks map
    37  			if v == nil {
    38  				delete(s.Hooks, k)
    39  			}
    40  		}
    41  	}
    42  }
    43  
    44  // Copy LocalState. Returns nil if nil.
    45  func (s *LocalState) Copy() *LocalState {
    46  	if s == nil {
    47  		return nil
    48  	}
    49  
    50  	// Create a copy
    51  	c := &LocalState{
    52  		Hooks:         make(map[string]*HookState, len(s.Hooks)),
    53  		DriverNetwork: s.DriverNetwork.Copy(),
    54  		TaskHandle:    s.TaskHandle.Copy(),
    55  	}
    56  
    57  	// Copy the hook state
    58  	for h, state := range s.Hooks {
    59  		c.Hooks[h] = state.Copy()
    60  	}
    61  
    62  	return c
    63  }
    64  
    65  type HookState struct {
    66  	// Prestart is true if the hook has run Prestart successfully and does
    67  	// not need to run again
    68  	PrestartDone bool
    69  
    70  	// Data allows hooks to persist arbitrary state.
    71  	Data map[string]string
    72  
    73  	// Environment variables set by the hook that will continue to be set
    74  	// even if PrestartDone=true.
    75  	Env map[string]string
    76  }
    77  
    78  // Copy HookState. Returns nil if its nil.
    79  func (h *HookState) Copy() *HookState {
    80  	if h == nil {
    81  		return nil
    82  	}
    83  
    84  	c := new(HookState)
    85  	*c = *h
    86  	c.Data = helper.CopyMapStringString(h.Data)
    87  	c.Env = helper.CopyMapStringString(h.Env)
    88  	return c
    89  }
    90  
    91  func (h *HookState) Equal(o *HookState) bool {
    92  	if h == nil || o == nil {
    93  		return h == o
    94  	}
    95  
    96  	if h.PrestartDone != o.PrestartDone {
    97  		return false
    98  	}
    99  
   100  	if !helper.CompareMapStringString(h.Data, o.Data) {
   101  		return false
   102  	}
   103  
   104  	return helper.CompareMapStringString(h.Env, o.Env)
   105  }