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  }