1package state
2
3import (
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.
10type 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
21func 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.
30func (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.
45func (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
65type 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.
79func (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
91func (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}
106