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  }