github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/state/db_test.go (about) 1 package state 2 3 import ( 4 "io/ioutil" 5 "os" 6 "reflect" 7 "sync" 8 "testing" 9 "time" 10 11 trstate "github.com/hashicorp/nomad/client/allocrunner/taskrunner/state" 12 dmstate "github.com/hashicorp/nomad/client/devicemanager/state" 13 "github.com/hashicorp/nomad/client/dynamicplugins" 14 driverstate "github.com/hashicorp/nomad/client/pluginmanager/drivermanager/state" 15 "github.com/hashicorp/nomad/helper/testlog" 16 "github.com/hashicorp/nomad/nomad/mock" 17 "github.com/hashicorp/nomad/nomad/structs" 18 "github.com/kr/pretty" 19 "github.com/stretchr/testify/require" 20 ) 21 22 func setupBoltStateDB(t *testing.T) (*BoltStateDB, func()) { 23 dir, err := ioutil.TempDir("", "nomadtest") 24 require.NoError(t, err) 25 26 db, err := NewBoltStateDB(testlog.HCLogger(t), dir) 27 if err != nil { 28 if err := os.RemoveAll(dir); err != nil { 29 t.Logf("error removing boltdb dir: %v", err) 30 } 31 t.Fatalf("error creating boltdb: %v", err) 32 } 33 34 cleanup := func() { 35 if err := db.Close(); err != nil { 36 t.Errorf("error closing boltdb: %v", err) 37 } 38 if err := os.RemoveAll(dir); err != nil { 39 t.Logf("error removing boltdb dir: %v", err) 40 } 41 } 42 43 return db.(*BoltStateDB), cleanup 44 } 45 46 func testDB(t *testing.T, f func(*testing.T, StateDB)) { 47 boltdb, cleanup := setupBoltStateDB(t) 48 defer cleanup() 49 50 memdb := NewMemDB(testlog.HCLogger(t)) 51 52 impls := []StateDB{boltdb, memdb} 53 54 for _, db := range impls { 55 db := db 56 t.Run(db.Name(), func(t *testing.T) { 57 f(t, db) 58 }) 59 } 60 } 61 62 // TestStateDB_Allocations asserts the behavior of GetAllAllocations, PutAllocation, and 63 // DeleteAllocationBucket for all operational StateDB implementations. 64 func TestStateDB_Allocations(t *testing.T) { 65 t.Parallel() 66 67 testDB(t, func(t *testing.T, db StateDB) { 68 require := require.New(t) 69 70 // Empty database should return empty non-nil results 71 allocs, errs, err := db.GetAllAllocations() 72 require.NoError(err) 73 require.NotNil(allocs) 74 require.Empty(allocs) 75 require.NotNil(errs) 76 require.Empty(errs) 77 78 // Put allocations 79 alloc1 := mock.Alloc() 80 alloc2 := mock.BatchAlloc() 81 82 require.NoError(db.PutAllocation(alloc1)) 83 require.NoError(db.PutAllocation(alloc2)) 84 85 // Retrieve them 86 allocs, errs, err = db.GetAllAllocations() 87 require.NoError(err) 88 require.NotNil(allocs) 89 require.Len(allocs, 2) 90 for _, a := range allocs { 91 switch a.ID { 92 case alloc1.ID: 93 if !reflect.DeepEqual(a, alloc1) { 94 pretty.Ldiff(t, a, alloc1) 95 t.Fatalf("alloc %q unequal", a.ID) 96 } 97 case alloc2.ID: 98 if !reflect.DeepEqual(a, alloc2) { 99 pretty.Ldiff(t, a, alloc2) 100 t.Fatalf("alloc %q unequal", a.ID) 101 } 102 default: 103 t.Fatalf("unexpected alloc id %q", a.ID) 104 } 105 } 106 require.NotNil(errs) 107 require.Empty(errs) 108 109 // Add another 110 alloc3 := mock.SystemAlloc() 111 require.NoError(db.PutAllocation(alloc3)) 112 allocs, errs, err = db.GetAllAllocations() 113 require.NoError(err) 114 require.NotNil(allocs) 115 require.Len(allocs, 3) 116 require.Contains(allocs, alloc1) 117 require.Contains(allocs, alloc2) 118 require.Contains(allocs, alloc3) 119 require.NotNil(errs) 120 require.Empty(errs) 121 122 // Deleting a nonexistent alloc is a noop 123 require.NoError(db.DeleteAllocationBucket("asdf")) 124 allocs, _, err = db.GetAllAllocations() 125 require.NoError(err) 126 require.NotNil(allocs) 127 require.Len(allocs, 3) 128 129 // Delete alloc1 130 require.NoError(db.DeleteAllocationBucket(alloc1.ID)) 131 allocs, errs, err = db.GetAllAllocations() 132 require.NoError(err) 133 require.NotNil(allocs) 134 require.Len(allocs, 2) 135 require.Contains(allocs, alloc2) 136 require.Contains(allocs, alloc3) 137 require.NotNil(errs) 138 require.Empty(errs) 139 }) 140 } 141 142 // Integer division, rounded up. 143 func ceilDiv(a, b int) int { 144 return (a + b - 1) / b 145 } 146 147 // TestStateDB_Batch asserts the behavior of PutAllocation, PutNetworkStatus and 148 // DeleteAllocationBucket in batch mode, for all operational StateDB implementations. 149 func TestStateDB_Batch(t *testing.T) { 150 t.Parallel() 151 152 testDB(t, func(t *testing.T, db StateDB) { 153 require := require.New(t) 154 155 // For BoltDB, get initial tx_id 156 var getTxID func() int 157 var prevTxID int 158 var batchDelay time.Duration 159 var batchSize int 160 if boltStateDB, ok := db.(*BoltStateDB); ok { 161 boltdb := boltStateDB.DB().BoltDB() 162 getTxID = func() int { 163 tx, err := boltdb.Begin(true) 164 require.NoError(err) 165 defer tx.Rollback() 166 return tx.ID() 167 } 168 prevTxID = getTxID() 169 batchDelay = boltdb.MaxBatchDelay 170 batchSize = boltdb.MaxBatchSize 171 } 172 173 // Write 1000 allocations and network statuses in batch mode 174 startTime := time.Now() 175 const numAllocs = 1000 176 var allocs []*structs.Allocation 177 for i := 0; i < numAllocs; i++ { 178 allocs = append(allocs, mock.Alloc()) 179 } 180 var wg sync.WaitGroup 181 for _, alloc := range allocs { 182 wg.Add(1) 183 go func(alloc *structs.Allocation) { 184 require.NoError(db.PutNetworkStatus(alloc.ID, mock.AllocNetworkStatus(), WithBatchMode())) 185 require.NoError(db.PutAllocation(alloc, WithBatchMode())) 186 wg.Done() 187 }(alloc) 188 } 189 wg.Wait() 190 191 // Check BoltDB actually combined PutAllocation calls into much fewer transactions. 192 // The actual number of transactions depends on how fast the goroutines are spawned, 193 // with every batchDelay (10ms by default) period saved in a separate transaction, 194 // plus each transaction is limited to batchSize writes (1000 by default). 195 // See boltdb MaxBatchDelay and MaxBatchSize parameters for more details. 196 if getTxID != nil { 197 numTransactions := getTxID() - prevTxID 198 writeTime := time.Now().Sub(startTime) 199 expectedNumTransactions := ceilDiv(2 * numAllocs, batchSize) + ceilDiv(int(writeTime), int(batchDelay)) 200 require.LessOrEqual(numTransactions, expectedNumTransactions) 201 prevTxID = getTxID() 202 } 203 204 // Retrieve allocs and make sure they are the same (order can differ) 205 readAllocs, errs, err := db.GetAllAllocations() 206 require.NoError(err) 207 require.NotNil(readAllocs) 208 require.Len(readAllocs, len(allocs)) 209 require.NotNil(errs) 210 require.Empty(errs) 211 212 readAllocsById := make(map[string]*structs.Allocation) 213 for _, readAlloc := range readAllocs { 214 readAllocsById[readAlloc.ID] = readAlloc 215 } 216 for _, alloc := range allocs { 217 readAlloc, ok := readAllocsById[alloc.ID] 218 if !ok { 219 t.Fatalf("no alloc with ID=%q", alloc.ID) 220 } 221 if !reflect.DeepEqual(readAlloc, alloc) { 222 pretty.Ldiff(t, readAlloc, alloc) 223 t.Fatalf("alloc %q unequal", alloc.ID) 224 } 225 } 226 227 // Delete all allocs in batch mode 228 startTime = time.Now() 229 for _, alloc := range allocs { 230 wg.Add(1) 231 go func(alloc *structs.Allocation) { 232 require.NoError(db.DeleteAllocationBucket(alloc.ID, WithBatchMode())) 233 wg.Done() 234 }(alloc) 235 } 236 wg.Wait() 237 238 // Check BoltDB combined DeleteAllocationBucket calls into much fewer transactions. 239 if getTxID != nil { 240 numTransactions := getTxID() - prevTxID 241 writeTime := time.Now().Sub(startTime) 242 expectedNumTransactions := ceilDiv(numAllocs, batchSize) + ceilDiv(int(writeTime), int(batchDelay)) 243 require.LessOrEqual(numTransactions, expectedNumTransactions) 244 prevTxID = getTxID() 245 } 246 247 // Check all allocs were deleted. 248 readAllocs, errs, err = db.GetAllAllocations() 249 require.NoError(err) 250 require.Empty(readAllocs) 251 require.Empty(errs) 252 }) 253 } 254 255 // TestStateDB_TaskState asserts the behavior of task state related StateDB 256 // methods. 257 func TestStateDB_TaskState(t *testing.T) { 258 t.Parallel() 259 260 testDB(t, func(t *testing.T, db StateDB) { 261 require := require.New(t) 262 263 // Getting nonexistent state should return nils 264 ls, ts, err := db.GetTaskRunnerState("allocid", "taskname") 265 require.NoError(err) 266 require.Nil(ls) 267 require.Nil(ts) 268 269 // Putting TaskState without first putting the allocation should work 270 state := structs.NewTaskState() 271 state.Failed = true // set a non-default value 272 require.NoError(db.PutTaskState("allocid", "taskname", state)) 273 274 // Getting should return the available state 275 ls, ts, err = db.GetTaskRunnerState("allocid", "taskname") 276 require.NoError(err) 277 require.Nil(ls) 278 require.Equal(state, ts) 279 280 // Deleting a nonexistent task should not error 281 require.NoError(db.DeleteTaskBucket("adsf", "asdf")) 282 require.NoError(db.DeleteTaskBucket("asllocid", "asdf")) 283 284 // Data should be untouched 285 ls, ts, err = db.GetTaskRunnerState("allocid", "taskname") 286 require.NoError(err) 287 require.Nil(ls) 288 require.Equal(state, ts) 289 290 // Deleting the task should remove the state 291 require.NoError(db.DeleteTaskBucket("allocid", "taskname")) 292 ls, ts, err = db.GetTaskRunnerState("allocid", "taskname") 293 require.NoError(err) 294 require.Nil(ls) 295 require.Nil(ts) 296 297 // Putting LocalState should work just like TaskState 298 origLocalState := trstate.NewLocalState() 299 require.NoError(db.PutTaskRunnerLocalState("allocid", "taskname", origLocalState)) 300 ls, ts, err = db.GetTaskRunnerState("allocid", "taskname") 301 require.NoError(err) 302 require.Equal(origLocalState, ls) 303 require.Nil(ts) 304 }) 305 } 306 307 // TestStateDB_DeviceManager asserts the behavior of device manager state related StateDB 308 // methods. 309 func TestStateDB_DeviceManager(t *testing.T) { 310 t.Parallel() 311 312 testDB(t, func(t *testing.T, db StateDB) { 313 require := require.New(t) 314 315 // Getting nonexistent state should return nils 316 ps, err := db.GetDevicePluginState() 317 require.NoError(err) 318 require.Nil(ps) 319 320 // Putting PluginState should work 321 state := &dmstate.PluginState{} 322 require.NoError(db.PutDevicePluginState(state)) 323 324 // Getting should return the available state 325 ps, err = db.GetDevicePluginState() 326 require.NoError(err) 327 require.NotNil(ps) 328 require.Equal(state, ps) 329 }) 330 } 331 332 // TestStateDB_DriverManager asserts the behavior of device manager state related StateDB 333 // methods. 334 func TestStateDB_DriverManager(t *testing.T) { 335 t.Parallel() 336 337 testDB(t, func(t *testing.T, db StateDB) { 338 require := require.New(t) 339 340 // Getting nonexistent state should return nils 341 ps, err := db.GetDriverPluginState() 342 require.NoError(err) 343 require.Nil(ps) 344 345 // Putting PluginState should work 346 state := &driverstate.PluginState{} 347 require.NoError(db.PutDriverPluginState(state)) 348 349 // Getting should return the available state 350 ps, err = db.GetDriverPluginState() 351 require.NoError(err) 352 require.NotNil(ps) 353 require.Equal(state, ps) 354 }) 355 } 356 357 // TestStateDB_DynamicRegistry asserts the behavior of dynamic registry state related StateDB 358 // methods. 359 func TestStateDB_DynamicRegistry(t *testing.T) { 360 t.Parallel() 361 362 testDB(t, func(t *testing.T, db StateDB) { 363 require := require.New(t) 364 365 // Getting nonexistent state should return nils 366 ps, err := db.GetDynamicPluginRegistryState() 367 require.NoError(err) 368 require.Nil(ps) 369 370 // Putting PluginState should work 371 state := &dynamicplugins.RegistryState{} 372 require.NoError(db.PutDynamicPluginRegistryState(state)) 373 374 // Getting should return the available state 375 ps, err = db.GetDynamicPluginRegistryState() 376 require.NoError(err) 377 require.NotNil(ps) 378 require.Equal(state, ps) 379 }) 380 } 381 382 // TestStateDB_Upgrade asserts calling Upgrade on new databases always 383 // succeeds. 384 func TestStateDB_Upgrade(t *testing.T) { 385 t.Parallel() 386 387 testDB(t, func(t *testing.T, db StateDB) { 388 require.NoError(t, db.Upgrade()) 389 }) 390 }