github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/state/memdb.go (about) 1 package state 2 3 import ( 4 "sync" 5 6 hclog "github.com/hashicorp/go-hclog" 7 "github.com/hashicorp/nomad/client/allocrunner/taskrunner/state" 8 dmstate "github.com/hashicorp/nomad/client/devicemanager/state" 9 "github.com/hashicorp/nomad/client/dynamicplugins" 10 driverstate "github.com/hashicorp/nomad/client/pluginmanager/drivermanager/state" 11 "github.com/hashicorp/nomad/nomad/structs" 12 ) 13 14 // MemDB implements a StateDB that stores data in memory and should only be 15 // used for testing. All methods are safe for concurrent use. 16 type MemDB struct { 17 // alloc_id -> value 18 allocs map[string]*structs.Allocation 19 20 // alloc_id -> value 21 deployStatus map[string]*structs.AllocDeploymentStatus 22 23 // alloc_id -> value 24 networkStatus map[string]*structs.AllocNetworkStatus 25 26 // alloc_id -> task_name -> value 27 localTaskState map[string]map[string]*state.LocalState 28 taskState map[string]map[string]*structs.TaskState 29 30 // devicemanager -> plugin-state 31 devManagerPs *dmstate.PluginState 32 33 // drivermanager -> plugin-state 34 driverManagerPs *driverstate.PluginState 35 36 // dynamicmanager -> registry-state 37 dynamicManagerPs *dynamicplugins.RegistryState 38 39 logger hclog.Logger 40 41 mu sync.RWMutex 42 } 43 44 func NewMemDB(logger hclog.Logger) *MemDB { 45 logger = logger.Named("memdb") 46 return &MemDB{ 47 allocs: make(map[string]*structs.Allocation), 48 deployStatus: make(map[string]*structs.AllocDeploymentStatus), 49 networkStatus: make(map[string]*structs.AllocNetworkStatus), 50 localTaskState: make(map[string]map[string]*state.LocalState), 51 taskState: make(map[string]map[string]*structs.TaskState), 52 logger: logger, 53 } 54 } 55 56 func (m *MemDB) Name() string { 57 return "memdb" 58 } 59 60 func (m *MemDB) Upgrade() error { 61 return nil 62 } 63 64 func (m *MemDB) GetAllAllocations() ([]*structs.Allocation, map[string]error, error) { 65 m.mu.RLock() 66 defer m.mu.RUnlock() 67 68 allocs := make([]*structs.Allocation, 0, len(m.allocs)) 69 for _, v := range m.allocs { 70 allocs = append(allocs, v) 71 } 72 73 return allocs, map[string]error{}, nil 74 } 75 76 func (m *MemDB) PutAllocation(alloc *structs.Allocation, opts ...WriteOption) error { 77 m.mu.Lock() 78 defer m.mu.Unlock() 79 m.allocs[alloc.ID] = alloc 80 return nil 81 } 82 83 func (m *MemDB) GetDeploymentStatus(allocID string) (*structs.AllocDeploymentStatus, error) { 84 m.mu.Lock() 85 defer m.mu.Unlock() 86 return m.deployStatus[allocID], nil 87 } 88 89 func (m *MemDB) PutDeploymentStatus(allocID string, ds *structs.AllocDeploymentStatus) error { 90 m.mu.Lock() 91 m.deployStatus[allocID] = ds 92 defer m.mu.Unlock() 93 return nil 94 } 95 96 func (m *MemDB) GetNetworkStatus(allocID string) (*structs.AllocNetworkStatus, error) { 97 m.mu.Lock() 98 defer m.mu.Unlock() 99 return m.networkStatus[allocID], nil 100 } 101 102 func (m *MemDB) PutNetworkStatus(allocID string, ns *structs.AllocNetworkStatus, opts ...WriteOption) error { 103 m.mu.Lock() 104 m.networkStatus[allocID] = ns 105 defer m.mu.Unlock() 106 return nil 107 } 108 109 func (m *MemDB) GetTaskRunnerState(allocID string, taskName string) (*state.LocalState, *structs.TaskState, error) { 110 m.mu.RLock() 111 defer m.mu.RUnlock() 112 113 var ls *state.LocalState 114 var ts *structs.TaskState 115 116 // Local Task State 117 allocLocalTS := m.localTaskState[allocID] 118 if len(allocLocalTS) != 0 { 119 ls = allocLocalTS[taskName] 120 } 121 122 // Task State 123 allocTS := m.taskState[allocID] 124 if len(allocTS) != 0 { 125 ts = allocTS[taskName] 126 } 127 128 return ls, ts, nil 129 } 130 131 func (m *MemDB) PutTaskRunnerLocalState(allocID string, taskName string, val *state.LocalState) error { 132 m.mu.Lock() 133 defer m.mu.Unlock() 134 135 if alts, ok := m.localTaskState[allocID]; ok { 136 alts[taskName] = val.Copy() 137 return nil 138 } 139 140 m.localTaskState[allocID] = map[string]*state.LocalState{ 141 taskName: val.Copy(), 142 } 143 144 return nil 145 } 146 147 func (m *MemDB) PutTaskState(allocID string, taskName string, state *structs.TaskState) error { 148 m.mu.Lock() 149 defer m.mu.Unlock() 150 151 if ats, ok := m.taskState[allocID]; ok { 152 ats[taskName] = state.Copy() 153 return nil 154 } 155 156 m.taskState[allocID] = map[string]*structs.TaskState{ 157 taskName: state.Copy(), 158 } 159 160 return nil 161 } 162 163 func (m *MemDB) DeleteTaskBucket(allocID, taskName string) error { 164 m.mu.Lock() 165 defer m.mu.Unlock() 166 167 if ats, ok := m.taskState[allocID]; ok { 168 delete(ats, taskName) 169 } 170 171 if alts, ok := m.localTaskState[allocID]; ok { 172 delete(alts, taskName) 173 } 174 175 return nil 176 } 177 178 func (m *MemDB) DeleteAllocationBucket(allocID string, opts ...WriteOption) error { 179 m.mu.Lock() 180 defer m.mu.Unlock() 181 182 delete(m.allocs, allocID) 183 delete(m.taskState, allocID) 184 delete(m.localTaskState, allocID) 185 186 return nil 187 } 188 189 func (m *MemDB) PutDevicePluginState(ps *dmstate.PluginState) error { 190 m.mu.Lock() 191 defer m.mu.Unlock() 192 m.devManagerPs = ps 193 return nil 194 } 195 196 // GetDevicePluginState stores the device manager's plugin state or returns an 197 // error. 198 func (m *MemDB) GetDevicePluginState() (*dmstate.PluginState, error) { 199 m.mu.Lock() 200 defer m.mu.Unlock() 201 return m.devManagerPs, nil 202 } 203 204 func (m *MemDB) GetDriverPluginState() (*driverstate.PluginState, error) { 205 m.mu.Lock() 206 defer m.mu.Unlock() 207 return m.driverManagerPs, nil 208 } 209 210 func (m *MemDB) PutDriverPluginState(ps *driverstate.PluginState) error { 211 m.mu.Lock() 212 defer m.mu.Unlock() 213 m.driverManagerPs = ps 214 return nil 215 } 216 217 func (m *MemDB) GetDynamicPluginRegistryState() (*dynamicplugins.RegistryState, error) { 218 m.mu.Lock() 219 defer m.mu.Unlock() 220 return m.dynamicManagerPs, nil 221 } 222 223 func (m *MemDB) PutDynamicPluginRegistryState(ps *dynamicplugins.RegistryState) error { 224 m.mu.Lock() 225 defer m.mu.Unlock() 226 m.dynamicManagerPs = ps 227 return nil 228 } 229 230 func (m *MemDB) Close() error { 231 m.mu.Lock() 232 defer m.mu.Unlock() 233 234 // Set everything to nil to blow up on further use 235 m.allocs = nil 236 m.taskState = nil 237 m.localTaskState = nil 238 239 return nil 240 }