github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/nomad/state/state_store_test.go (about)

     1  package state
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"reflect"
     7  	"sort"
     8  	"strings"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/hashicorp/go-memdb"
    13  	"github.com/kr/pretty"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  
    17  	"github.com/hashicorp/nomad/helper"
    18  	"github.com/hashicorp/nomad/helper/uuid"
    19  	"github.com/hashicorp/nomad/nomad/mock"
    20  	"github.com/hashicorp/nomad/nomad/structs"
    21  )
    22  
    23  func testStateStore(t *testing.T) *StateStore {
    24  	return TestStateStore(t)
    25  }
    26  
    27  func TestStateStore_Blocking_Error(t *testing.T) {
    28  	t.Parallel()
    29  
    30  	expected := fmt.Errorf("test error")
    31  	errFn := func(memdb.WatchSet, *StateStore) (interface{}, uint64, error) {
    32  		return nil, 0, expected
    33  	}
    34  
    35  	state := testStateStore(t)
    36  	_, idx, err := state.BlockingQuery(errFn, 10, context.Background())
    37  	assert.EqualError(t, err, expected.Error())
    38  	assert.Zero(t, idx)
    39  }
    40  
    41  func TestStateStore_Blocking_Timeout(t *testing.T) {
    42  	t.Parallel()
    43  
    44  	noopFn := func(memdb.WatchSet, *StateStore) (interface{}, uint64, error) {
    45  		return nil, 5, nil
    46  	}
    47  
    48  	state := testStateStore(t)
    49  	timeout := time.Now().Add(250 * time.Millisecond)
    50  	deadlineCtx, cancel := context.WithDeadline(context.Background(), timeout)
    51  	defer cancel()
    52  
    53  	_, idx, err := state.BlockingQuery(noopFn, 10, deadlineCtx)
    54  	assert.EqualError(t, err, context.DeadlineExceeded.Error())
    55  	assert.EqualValues(t, 5, idx)
    56  	assert.WithinDuration(t, timeout, time.Now(), 100*time.Millisecond)
    57  }
    58  
    59  func TestStateStore_Blocking_MinQuery(t *testing.T) {
    60  	t.Parallel()
    61  
    62  	node := mock.Node()
    63  	count := 0
    64  	queryFn := func(ws memdb.WatchSet, s *StateStore) (interface{}, uint64, error) {
    65  		_, err := s.NodeByID(ws, node.ID)
    66  		if err != nil {
    67  			return nil, 0, err
    68  		}
    69  
    70  		count++
    71  		if count == 1 {
    72  			return false, 5, nil
    73  		} else if count > 2 {
    74  			return false, 20, fmt.Errorf("called too many times")
    75  		}
    76  
    77  		return true, 11, nil
    78  	}
    79  
    80  	state := testStateStore(t)
    81  	timeout := time.Now().Add(100 * time.Millisecond)
    82  	deadlineCtx, cancel := context.WithDeadline(context.Background(), timeout)
    83  	defer cancel()
    84  
    85  	time.AfterFunc(5*time.Millisecond, func() {
    86  		state.UpsertNode(structs.MsgTypeTestSetup, 11, node)
    87  	})
    88  
    89  	resp, idx, err := state.BlockingQuery(queryFn, 10, deadlineCtx)
    90  	if assert.Nil(t, err) {
    91  		assert.Equal(t, 2, count)
    92  		assert.EqualValues(t, 11, idx)
    93  		assert.True(t, resp.(bool))
    94  	}
    95  }
    96  
    97  // COMPAT 0.11: Uses AllocUpdateRequest.Alloc
    98  // This test checks that:
    99  // 1) The job is denormalized
   100  // 2) Allocations are created
   101  func TestStateStore_UpsertPlanResults_AllocationsCreated_Denormalized(t *testing.T) {
   102  	t.Parallel()
   103  
   104  	state := testStateStore(t)
   105  	alloc := mock.Alloc()
   106  	job := alloc.Job
   107  	alloc.Job = nil
   108  
   109  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, job); err != nil {
   110  		t.Fatalf("err: %v", err)
   111  	}
   112  
   113  	eval := mock.Eval()
   114  	eval.JobID = job.ID
   115  
   116  	// Create an eval
   117  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}); err != nil {
   118  		t.Fatalf("err: %v", err)
   119  	}
   120  
   121  	// Create a plan result
   122  	res := structs.ApplyPlanResultsRequest{
   123  		AllocUpdateRequest: structs.AllocUpdateRequest{
   124  			Alloc: []*structs.Allocation{alloc},
   125  			Job:   job,
   126  		},
   127  		EvalID: eval.ID,
   128  	}
   129  	assert := assert.New(t)
   130  	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
   131  	assert.Nil(err)
   132  
   133  	ws := memdb.NewWatchSet()
   134  	out, err := state.AllocByID(ws, alloc.ID)
   135  	assert.Nil(err)
   136  	assert.Equal(alloc, out)
   137  
   138  	index, err := state.Index("allocs")
   139  	assert.Nil(err)
   140  	assert.EqualValues(1000, index)
   141  
   142  	if watchFired(ws) {
   143  		t.Fatalf("bad")
   144  	}
   145  
   146  	evalOut, err := state.EvalByID(ws, eval.ID)
   147  	assert.Nil(err)
   148  	assert.NotNil(evalOut)
   149  	assert.EqualValues(1000, evalOut.ModifyIndex)
   150  }
   151  
   152  // This test checks that:
   153  // 1) The job is denormalized
   154  // 2) Allocations are denormalized and updated with the diff
   155  // That stopped allocs Job is unmodified
   156  func TestStateStore_UpsertPlanResults_AllocationsDenormalized(t *testing.T) {
   157  	t.Parallel()
   158  
   159  	state := testStateStore(t)
   160  	alloc := mock.Alloc()
   161  	job := alloc.Job
   162  	alloc.Job = nil
   163  
   164  	stoppedAlloc := mock.Alloc()
   165  	stoppedAlloc.Job = job
   166  	stoppedAllocDiff := &structs.AllocationDiff{
   167  		ID:                 stoppedAlloc.ID,
   168  		DesiredDescription: "desired desc",
   169  		ClientStatus:       structs.AllocClientStatusLost,
   170  	}
   171  	preemptedAlloc := mock.Alloc()
   172  	preemptedAlloc.Job = job
   173  	preemptedAllocDiff := &structs.AllocationDiff{
   174  		ID:                    preemptedAlloc.ID,
   175  		PreemptedByAllocation: alloc.ID,
   176  	}
   177  
   178  	require := require.New(t)
   179  	require.NoError(state.UpsertAllocs(structs.MsgTypeTestSetup, 900, []*structs.Allocation{stoppedAlloc, preemptedAlloc}))
   180  	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 999, job))
   181  
   182  	// modify job and ensure that stopped and preempted alloc point to original Job
   183  	mJob := job.Copy()
   184  	mJob.TaskGroups[0].Name = "other"
   185  
   186  	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 1001, mJob))
   187  
   188  	eval := mock.Eval()
   189  	eval.JobID = job.ID
   190  
   191  	// Create an eval
   192  	require.NoError(state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}))
   193  
   194  	// Create a plan result
   195  	res := structs.ApplyPlanResultsRequest{
   196  		AllocUpdateRequest: structs.AllocUpdateRequest{
   197  			AllocsUpdated: []*structs.Allocation{alloc},
   198  			AllocsStopped: []*structs.AllocationDiff{stoppedAllocDiff},
   199  			Job:           mJob,
   200  		},
   201  		EvalID:          eval.ID,
   202  		AllocsPreempted: []*structs.AllocationDiff{preemptedAllocDiff},
   203  	}
   204  	assert := assert.New(t)
   205  	planModifyIndex := uint64(1000)
   206  	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, planModifyIndex, &res)
   207  	require.NoError(err)
   208  
   209  	ws := memdb.NewWatchSet()
   210  	out, err := state.AllocByID(ws, alloc.ID)
   211  	require.NoError(err)
   212  	assert.Equal(alloc, out)
   213  
   214  	outJob, err := state.JobByID(ws, job.Namespace, job.ID)
   215  	require.NoError(err)
   216  	require.Equal(mJob.TaskGroups, outJob.TaskGroups)
   217  	require.NotEmpty(job.TaskGroups, outJob.TaskGroups)
   218  
   219  	updatedStoppedAlloc, err := state.AllocByID(ws, stoppedAlloc.ID)
   220  	require.NoError(err)
   221  	assert.Equal(stoppedAllocDiff.DesiredDescription, updatedStoppedAlloc.DesiredDescription)
   222  	assert.Equal(structs.AllocDesiredStatusStop, updatedStoppedAlloc.DesiredStatus)
   223  	assert.Equal(stoppedAllocDiff.ClientStatus, updatedStoppedAlloc.ClientStatus)
   224  	assert.Equal(planModifyIndex, updatedStoppedAlloc.AllocModifyIndex)
   225  	assert.Equal(planModifyIndex, updatedStoppedAlloc.AllocModifyIndex)
   226  	assert.Equal(job.TaskGroups, updatedStoppedAlloc.Job.TaskGroups)
   227  
   228  	updatedPreemptedAlloc, err := state.AllocByID(ws, preemptedAlloc.ID)
   229  	require.NoError(err)
   230  	assert.Equal(structs.AllocDesiredStatusEvict, updatedPreemptedAlloc.DesiredStatus)
   231  	assert.Equal(preemptedAllocDiff.PreemptedByAllocation, updatedPreemptedAlloc.PreemptedByAllocation)
   232  	assert.Equal(planModifyIndex, updatedPreemptedAlloc.AllocModifyIndex)
   233  	assert.Equal(planModifyIndex, updatedPreemptedAlloc.AllocModifyIndex)
   234  	assert.Equal(job.TaskGroups, updatedPreemptedAlloc.Job.TaskGroups)
   235  
   236  	index, err := state.Index("allocs")
   237  	require.NoError(err)
   238  	assert.EqualValues(planModifyIndex, index)
   239  
   240  	require.False(watchFired(ws))
   241  
   242  	evalOut, err := state.EvalByID(ws, eval.ID)
   243  	require.NoError(err)
   244  	require.NotNil(evalOut)
   245  	assert.EqualValues(planModifyIndex, evalOut.ModifyIndex)
   246  
   247  }
   248  
   249  // This test checks that the deployment is created and allocations count towards
   250  // the deployment
   251  func TestStateStore_UpsertPlanResults_Deployment(t *testing.T) {
   252  	t.Parallel()
   253  
   254  	state := testStateStore(t)
   255  	alloc := mock.Alloc()
   256  	alloc2 := mock.Alloc()
   257  	job := alloc.Job
   258  	alloc.Job = nil
   259  	alloc2.Job = nil
   260  
   261  	d := mock.Deployment()
   262  	alloc.DeploymentID = d.ID
   263  	alloc2.DeploymentID = d.ID
   264  
   265  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, job); err != nil {
   266  		t.Fatalf("err: %v", err)
   267  	}
   268  
   269  	eval := mock.Eval()
   270  	eval.JobID = job.ID
   271  
   272  	// Create an eval
   273  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}); err != nil {
   274  		t.Fatalf("err: %v", err)
   275  	}
   276  
   277  	// Create a plan result
   278  	res := structs.ApplyPlanResultsRequest{
   279  		AllocUpdateRequest: structs.AllocUpdateRequest{
   280  			Alloc: []*structs.Allocation{alloc, alloc2},
   281  			Job:   job,
   282  		},
   283  		Deployment: d,
   284  		EvalID:     eval.ID,
   285  	}
   286  
   287  	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
   288  	if err != nil {
   289  		t.Fatalf("err: %v", err)
   290  	}
   291  
   292  	ws := memdb.NewWatchSet()
   293  	assert := assert.New(t)
   294  	out, err := state.AllocByID(ws, alloc.ID)
   295  	assert.Nil(err)
   296  	assert.Equal(alloc, out)
   297  
   298  	dout, err := state.DeploymentByID(ws, d.ID)
   299  	assert.Nil(err)
   300  	assert.NotNil(dout)
   301  
   302  	tg, ok := dout.TaskGroups[alloc.TaskGroup]
   303  	assert.True(ok)
   304  	assert.NotNil(tg)
   305  	assert.Equal(2, tg.PlacedAllocs)
   306  
   307  	evalOut, err := state.EvalByID(ws, eval.ID)
   308  	assert.Nil(err)
   309  	assert.NotNil(evalOut)
   310  	assert.EqualValues(1000, evalOut.ModifyIndex)
   311  
   312  	if watchFired(ws) {
   313  		t.Fatalf("bad")
   314  	}
   315  
   316  	// Update the allocs to be part of a new deployment
   317  	d2 := d.Copy()
   318  	d2.ID = uuid.Generate()
   319  
   320  	allocNew := alloc.Copy()
   321  	allocNew.DeploymentID = d2.ID
   322  	allocNew2 := alloc2.Copy()
   323  	allocNew2.DeploymentID = d2.ID
   324  
   325  	// Create another plan
   326  	res = structs.ApplyPlanResultsRequest{
   327  		AllocUpdateRequest: structs.AllocUpdateRequest{
   328  			Alloc: []*structs.Allocation{allocNew, allocNew2},
   329  			Job:   job,
   330  		},
   331  		Deployment: d2,
   332  		EvalID:     eval.ID,
   333  	}
   334  
   335  	err = state.UpsertPlanResults(structs.MsgTypeTestSetup, 1001, &res)
   336  	if err != nil {
   337  		t.Fatalf("err: %v", err)
   338  	}
   339  
   340  	dout, err = state.DeploymentByID(ws, d2.ID)
   341  	assert.Nil(err)
   342  	assert.NotNil(dout)
   343  
   344  	tg, ok = dout.TaskGroups[alloc.TaskGroup]
   345  	assert.True(ok)
   346  	assert.NotNil(tg)
   347  	assert.Equal(2, tg.PlacedAllocs)
   348  
   349  	evalOut, err = state.EvalByID(ws, eval.ID)
   350  	assert.Nil(err)
   351  	assert.NotNil(evalOut)
   352  	assert.EqualValues(1001, evalOut.ModifyIndex)
   353  }
   354  
   355  // This test checks that:
   356  // 1) Preempted allocations in plan results are updated
   357  // 2) Evals are inserted for preempted jobs
   358  func TestStateStore_UpsertPlanResults_PreemptedAllocs(t *testing.T) {
   359  	t.Parallel()
   360  	require := require.New(t)
   361  
   362  	state := testStateStore(t)
   363  	alloc := mock.Alloc()
   364  	job := alloc.Job
   365  	alloc.Job = nil
   366  
   367  	// Insert job
   368  	err := state.UpsertJob(structs.MsgTypeTestSetup, 999, job)
   369  	require.NoError(err)
   370  
   371  	// Create an eval
   372  	eval := mock.Eval()
   373  	eval.JobID = job.ID
   374  	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval})
   375  	require.NoError(err)
   376  
   377  	// Insert alloc that will be preempted in the plan
   378  	preemptedAlloc := mock.Alloc()
   379  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 2, []*structs.Allocation{preemptedAlloc})
   380  	require.NoError(err)
   381  
   382  	minimalPreemptedAlloc := &structs.Allocation{
   383  		ID:                    preemptedAlloc.ID,
   384  		PreemptedByAllocation: alloc.ID,
   385  		ModifyTime:            time.Now().Unix(),
   386  	}
   387  
   388  	// Create eval for preempted job
   389  	eval2 := mock.Eval()
   390  	eval2.JobID = preemptedAlloc.JobID
   391  
   392  	// Create a plan result
   393  	res := structs.ApplyPlanResultsRequest{
   394  		AllocUpdateRequest: structs.AllocUpdateRequest{
   395  			Alloc: []*structs.Allocation{alloc},
   396  			Job:   job,
   397  		},
   398  		EvalID:          eval.ID,
   399  		NodePreemptions: []*structs.Allocation{minimalPreemptedAlloc},
   400  		PreemptionEvals: []*structs.Evaluation{eval2},
   401  	}
   402  
   403  	err = state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
   404  	require.NoError(err)
   405  
   406  	ws := memdb.NewWatchSet()
   407  
   408  	// Verify alloc and eval created by plan
   409  	out, err := state.AllocByID(ws, alloc.ID)
   410  	require.NoError(err)
   411  	require.Equal(alloc, out)
   412  
   413  	index, err := state.Index("allocs")
   414  	require.NoError(err)
   415  	require.EqualValues(1000, index)
   416  
   417  	evalOut, err := state.EvalByID(ws, eval.ID)
   418  	require.NoError(err)
   419  	require.NotNil(evalOut)
   420  	require.EqualValues(1000, evalOut.ModifyIndex)
   421  
   422  	// Verify preempted alloc
   423  	preempted, err := state.AllocByID(ws, preemptedAlloc.ID)
   424  	require.NoError(err)
   425  	require.Equal(preempted.DesiredStatus, structs.AllocDesiredStatusEvict)
   426  	require.Equal(preempted.DesiredDescription, fmt.Sprintf("Preempted by alloc ID %v", alloc.ID))
   427  	require.Equal(preempted.Job.ID, preemptedAlloc.Job.ID)
   428  	require.Equal(preempted.Job, preemptedAlloc.Job)
   429  
   430  	// Verify eval for preempted job
   431  	preemptedJobEval, err := state.EvalByID(ws, eval2.ID)
   432  	require.NoError(err)
   433  	require.NotNil(preemptedJobEval)
   434  	require.EqualValues(1000, preemptedJobEval.ModifyIndex)
   435  
   436  }
   437  
   438  // This test checks that deployment updates are applied correctly
   439  func TestStateStore_UpsertPlanResults_DeploymentUpdates(t *testing.T) {
   440  	t.Parallel()
   441  	state := testStateStore(t)
   442  
   443  	// Create a job that applies to all
   444  	job := mock.Job()
   445  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, job); err != nil {
   446  		t.Fatalf("err: %v", err)
   447  	}
   448  
   449  	// Create a deployment that we will update its status
   450  	doutstanding := mock.Deployment()
   451  	doutstanding.JobID = job.ID
   452  
   453  	if err := state.UpsertDeployment(1000, doutstanding); err != nil {
   454  		t.Fatalf("err: %v", err)
   455  	}
   456  
   457  	eval := mock.Eval()
   458  	eval.JobID = job.ID
   459  
   460  	// Create an eval
   461  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1, []*structs.Evaluation{eval}); err != nil {
   462  		t.Fatalf("err: %v", err)
   463  	}
   464  	alloc := mock.Alloc()
   465  	alloc.Job = nil
   466  
   467  	dnew := mock.Deployment()
   468  	dnew.JobID = job.ID
   469  	alloc.DeploymentID = dnew.ID
   470  
   471  	// Update the old deployment
   472  	update := &structs.DeploymentStatusUpdate{
   473  		DeploymentID:      doutstanding.ID,
   474  		Status:            "foo",
   475  		StatusDescription: "bar",
   476  	}
   477  
   478  	// Create a plan result
   479  	res := structs.ApplyPlanResultsRequest{
   480  		AllocUpdateRequest: structs.AllocUpdateRequest{
   481  			Alloc: []*structs.Allocation{alloc},
   482  			Job:   job,
   483  		},
   484  		Deployment:        dnew,
   485  		DeploymentUpdates: []*structs.DeploymentStatusUpdate{update},
   486  		EvalID:            eval.ID,
   487  	}
   488  
   489  	err := state.UpsertPlanResults(structs.MsgTypeTestSetup, 1000, &res)
   490  	if err != nil {
   491  		t.Fatalf("err: %v", err)
   492  	}
   493  	assert := assert.New(t)
   494  	ws := memdb.NewWatchSet()
   495  
   496  	// Check the deployments are correctly updated.
   497  	dout, err := state.DeploymentByID(ws, dnew.ID)
   498  	assert.Nil(err)
   499  	assert.NotNil(dout)
   500  
   501  	tg, ok := dout.TaskGroups[alloc.TaskGroup]
   502  	assert.True(ok)
   503  	assert.NotNil(tg)
   504  	assert.Equal(1, tg.PlacedAllocs)
   505  
   506  	doutstandingout, err := state.DeploymentByID(ws, doutstanding.ID)
   507  	assert.Nil(err)
   508  	assert.NotNil(doutstandingout)
   509  	assert.Equal(update.Status, doutstandingout.Status)
   510  	assert.Equal(update.StatusDescription, doutstandingout.StatusDescription)
   511  	assert.EqualValues(1000, doutstandingout.ModifyIndex)
   512  
   513  	evalOut, err := state.EvalByID(ws, eval.ID)
   514  	assert.Nil(err)
   515  	assert.NotNil(evalOut)
   516  	assert.EqualValues(1000, evalOut.ModifyIndex)
   517  	if watchFired(ws) {
   518  		t.Fatalf("bad")
   519  	}
   520  }
   521  
   522  func TestStateStore_UpsertDeployment(t *testing.T) {
   523  	t.Parallel()
   524  
   525  	state := testStateStore(t)
   526  	deployment := mock.Deployment()
   527  
   528  	// Create a watchset so we can test that upsert fires the watch
   529  	ws := memdb.NewWatchSet()
   530  	_, err := state.DeploymentsByJobID(ws, deployment.Namespace, deployment.ID, true)
   531  	if err != nil {
   532  		t.Fatalf("bad: %v", err)
   533  	}
   534  
   535  	err = state.UpsertDeployment(1000, deployment)
   536  	if err != nil {
   537  		t.Fatalf("err: %v", err)
   538  	}
   539  	if !watchFired(ws) {
   540  		t.Fatalf("bad")
   541  	}
   542  
   543  	ws = memdb.NewWatchSet()
   544  	out, err := state.DeploymentByID(ws, deployment.ID)
   545  	if err != nil {
   546  		t.Fatalf("err: %v", err)
   547  	}
   548  
   549  	if !reflect.DeepEqual(deployment, out) {
   550  		t.Fatalf("bad: %#v %#v", deployment, out)
   551  	}
   552  
   553  	index, err := state.Index("deployment")
   554  	if err != nil {
   555  		t.Fatalf("err: %v", err)
   556  	}
   557  	if index != 1000 {
   558  		t.Fatalf("bad: %d", index)
   559  	}
   560  
   561  	if watchFired(ws) {
   562  		t.Fatalf("bad")
   563  	}
   564  }
   565  
   566  // Tests that deployments of older create index and same job id are not returned
   567  func TestStateStore_OldDeployment(t *testing.T) {
   568  	t.Parallel()
   569  
   570  	state := testStateStore(t)
   571  	job := mock.Job()
   572  	job.ID = "job1"
   573  	state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
   574  
   575  	deploy1 := mock.Deployment()
   576  	deploy1.JobID = job.ID
   577  	deploy1.JobCreateIndex = job.CreateIndex
   578  
   579  	deploy2 := mock.Deployment()
   580  	deploy2.JobID = job.ID
   581  	deploy2.JobCreateIndex = 11
   582  
   583  	require := require.New(t)
   584  
   585  	// Insert both deployments
   586  	err := state.UpsertDeployment(1001, deploy1)
   587  	require.Nil(err)
   588  
   589  	err = state.UpsertDeployment(1002, deploy2)
   590  	require.Nil(err)
   591  
   592  	ws := memdb.NewWatchSet()
   593  	// Should return both deployments
   594  	deploys, err := state.DeploymentsByJobID(ws, deploy1.Namespace, job.ID, true)
   595  	require.Nil(err)
   596  	require.Len(deploys, 2)
   597  
   598  	// Should only return deploy1
   599  	deploys, err = state.DeploymentsByJobID(ws, deploy1.Namespace, job.ID, false)
   600  	require.Nil(err)
   601  	require.Len(deploys, 1)
   602  	require.Equal(deploy1.ID, deploys[0].ID)
   603  }
   604  
   605  func TestStateStore_DeleteDeployment(t *testing.T) {
   606  	t.Parallel()
   607  
   608  	state := testStateStore(t)
   609  	d1 := mock.Deployment()
   610  	d2 := mock.Deployment()
   611  
   612  	err := state.UpsertDeployment(1000, d1)
   613  	if err != nil {
   614  		t.Fatalf("err: %v", err)
   615  	}
   616  	if err := state.UpsertDeployment(1001, d2); err != nil {
   617  		t.Fatalf("err: %v", err)
   618  	}
   619  
   620  	// Create a watchset so we can test that delete fires the watch
   621  	ws := memdb.NewWatchSet()
   622  	if _, err := state.DeploymentByID(ws, d1.ID); err != nil {
   623  		t.Fatalf("bad: %v", err)
   624  	}
   625  
   626  	err = state.DeleteDeployment(1002, []string{d1.ID, d2.ID})
   627  	if err != nil {
   628  		t.Fatalf("err: %v", err)
   629  	}
   630  
   631  	if !watchFired(ws) {
   632  		t.Fatalf("bad")
   633  	}
   634  
   635  	ws = memdb.NewWatchSet()
   636  	out, err := state.DeploymentByID(ws, d1.ID)
   637  	if err != nil {
   638  		t.Fatalf("err: %v", err)
   639  	}
   640  
   641  	if out != nil {
   642  		t.Fatalf("bad: %#v %#v", d1, out)
   643  	}
   644  
   645  	index, err := state.Index("deployment")
   646  	if err != nil {
   647  		t.Fatalf("err: %v", err)
   648  	}
   649  	if index != 1002 {
   650  		t.Fatalf("bad: %d", index)
   651  	}
   652  
   653  	if watchFired(ws) {
   654  		t.Fatalf("bad")
   655  	}
   656  }
   657  
   658  func TestStateStore_Deployments(t *testing.T) {
   659  	t.Parallel()
   660  
   661  	state := testStateStore(t)
   662  	var deployments []*structs.Deployment
   663  
   664  	for i := 0; i < 10; i++ {
   665  		deployment := mock.Deployment()
   666  		deployments = append(deployments, deployment)
   667  
   668  		err := state.UpsertDeployment(1000+uint64(i), deployment)
   669  		if err != nil {
   670  			t.Fatalf("err: %v", err)
   671  		}
   672  	}
   673  
   674  	ws := memdb.NewWatchSet()
   675  	iter, err := state.Deployments(ws)
   676  	if err != nil {
   677  		t.Fatalf("err: %v", err)
   678  	}
   679  
   680  	var out []*structs.Deployment
   681  	for {
   682  		raw := iter.Next()
   683  		if raw == nil {
   684  			break
   685  		}
   686  		out = append(out, raw.(*structs.Deployment))
   687  	}
   688  
   689  	lessThan := func(i, j int) bool {
   690  		return deployments[i].ID < deployments[j].ID
   691  	}
   692  	sort.Slice(deployments, lessThan)
   693  	sort.Slice(out, lessThan)
   694  
   695  	if !reflect.DeepEqual(deployments, out) {
   696  		t.Fatalf("bad: %#v %#v", deployments, out)
   697  	}
   698  
   699  	if watchFired(ws) {
   700  		t.Fatalf("bad")
   701  	}
   702  }
   703  
   704  func TestStateStore_DeploymentsByIDPrefix(t *testing.T) {
   705  	t.Parallel()
   706  
   707  	state := testStateStore(t)
   708  	deploy := mock.Deployment()
   709  
   710  	deploy.ID = "11111111-662e-d0ab-d1c9-3e434af7bdb4"
   711  	err := state.UpsertDeployment(1000, deploy)
   712  	if err != nil {
   713  		t.Fatalf("err: %v", err)
   714  	}
   715  
   716  	// Create a watchset so we can test that getters don't cause it to fire
   717  	ws := memdb.NewWatchSet()
   718  	iter, err := state.DeploymentsByIDPrefix(ws, deploy.Namespace, deploy.ID)
   719  	if err != nil {
   720  		t.Fatalf("err: %v", err)
   721  	}
   722  
   723  	gatherDeploys := func(iter memdb.ResultIterator) []*structs.Deployment {
   724  		var deploys []*structs.Deployment
   725  		for {
   726  			raw := iter.Next()
   727  			if raw == nil {
   728  				break
   729  			}
   730  			deploy := raw.(*structs.Deployment)
   731  			deploys = append(deploys, deploy)
   732  		}
   733  		return deploys
   734  	}
   735  
   736  	deploys := gatherDeploys(iter)
   737  	if len(deploys) != 1 {
   738  		t.Fatalf("err: %v", err)
   739  	}
   740  
   741  	if watchFired(ws) {
   742  		t.Fatalf("bad")
   743  	}
   744  
   745  	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "11")
   746  	if err != nil {
   747  		t.Fatalf("err: %v", err)
   748  	}
   749  
   750  	deploys = gatherDeploys(iter)
   751  	if len(deploys) != 1 {
   752  		t.Fatalf("err: %v", err)
   753  	}
   754  
   755  	deploy = mock.Deployment()
   756  	deploy.ID = "11222222-662e-d0ab-d1c9-3e434af7bdb4"
   757  	err = state.UpsertDeployment(1001, deploy)
   758  	if err != nil {
   759  		t.Fatalf("err: %v", err)
   760  	}
   761  
   762  	if !watchFired(ws) {
   763  		t.Fatalf("bad")
   764  	}
   765  
   766  	ws = memdb.NewWatchSet()
   767  	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "11")
   768  	if err != nil {
   769  		t.Fatalf("err: %v", err)
   770  	}
   771  
   772  	deploys = gatherDeploys(iter)
   773  	if len(deploys) != 2 {
   774  		t.Fatalf("err: %v", err)
   775  	}
   776  
   777  	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "1111")
   778  	if err != nil {
   779  		t.Fatalf("err: %v", err)
   780  	}
   781  
   782  	deploys = gatherDeploys(iter)
   783  	if len(deploys) != 1 {
   784  		t.Fatalf("err: %v", err)
   785  	}
   786  
   787  	if watchFired(ws) {
   788  		t.Fatalf("bad")
   789  	}
   790  }
   791  
   792  func TestStateStore_UpsertNode_Node(t *testing.T) {
   793  	t.Parallel()
   794  
   795  	require := require.New(t)
   796  	state := testStateStore(t)
   797  	node := mock.Node()
   798  
   799  	// Create a watchset so we can test that upsert fires the watch
   800  	ws := memdb.NewWatchSet()
   801  	_, err := state.NodeByID(ws, node.ID)
   802  	require.NoError(err)
   803  
   804  	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 1000, node))
   805  	require.True(watchFired(ws))
   806  
   807  	ws = memdb.NewWatchSet()
   808  	out, err := state.NodeByID(ws, node.ID)
   809  	require.NoError(err)
   810  
   811  	out2, err := state.NodeBySecretID(ws, node.SecretID)
   812  	require.NoError(err)
   813  	require.EqualValues(node, out)
   814  	require.EqualValues(node, out2)
   815  	require.Len(out.Events, 1)
   816  	require.Equal(NodeRegisterEventRegistered, out.Events[0].Message)
   817  
   818  	index, err := state.Index("nodes")
   819  	require.NoError(err)
   820  	require.EqualValues(1000, index)
   821  	require.False(watchFired(ws))
   822  
   823  	// Transition the node to down and then up and ensure we get a re-register
   824  	// event
   825  	down := out.Copy()
   826  	down.Status = structs.NodeStatusDown
   827  	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 1001, down))
   828  	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 1002, out))
   829  
   830  	out, err = state.NodeByID(ws, node.ID)
   831  	require.NoError(err)
   832  	require.Len(out.Events, 2)
   833  	require.Equal(NodeRegisterEventReregistered, out.Events[1].Message)
   834  }
   835  
   836  func TestStateStore_DeleteNode_Node(t *testing.T) {
   837  	t.Parallel()
   838  
   839  	state := testStateStore(t)
   840  
   841  	// Create and insert two nodes, which we'll delete
   842  	node0 := mock.Node()
   843  	node1 := mock.Node()
   844  	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node0)
   845  	require.NoError(t, err)
   846  	err = state.UpsertNode(structs.MsgTypeTestSetup, 1001, node1)
   847  	require.NoError(t, err)
   848  
   849  	// Create a watchset so we can test that delete fires the watch
   850  	ws := memdb.NewWatchSet()
   851  
   852  	// Check that both nodes are not nil
   853  	out, err := state.NodeByID(ws, node0.ID)
   854  	require.NoError(t, err)
   855  	require.NotNil(t, out)
   856  	out, err = state.NodeByID(ws, node1.ID)
   857  	require.NoError(t, err)
   858  	require.NotNil(t, out)
   859  
   860  	// Delete both nodes in a batch, fires the watch
   861  	err = state.DeleteNode(structs.MsgTypeTestSetup, 1002, []string{node0.ID, node1.ID})
   862  	require.NoError(t, err)
   863  	require.True(t, watchFired(ws))
   864  
   865  	// Check that both nodes are nil
   866  	ws = memdb.NewWatchSet()
   867  	out, err = state.NodeByID(ws, node0.ID)
   868  	require.NoError(t, err)
   869  	require.Nil(t, out)
   870  	out, err = state.NodeByID(ws, node1.ID)
   871  	require.NoError(t, err)
   872  	require.Nil(t, out)
   873  
   874  	// Ensure that the index is still at 1002, from DeleteNode
   875  	index, err := state.Index("nodes")
   876  	require.NoError(t, err)
   877  	require.Equal(t, uint64(1002), index)
   878  	require.False(t, watchFired(ws))
   879  }
   880  
   881  func TestStateStore_UpdateNodeStatus_Node(t *testing.T) {
   882  	t.Parallel()
   883  	require := require.New(t)
   884  
   885  	state := testStateStore(t)
   886  	node := mock.Node()
   887  
   888  	require.NoError(state.UpsertNode(structs.MsgTypeTestSetup, 800, node))
   889  
   890  	// Create a watchset so we can test that update node status fires the watch
   891  	ws := memdb.NewWatchSet()
   892  	_, err := state.NodeByID(ws, node.ID)
   893  	require.NoError(err)
   894  
   895  	event := &structs.NodeEvent{
   896  		Message:   "Node ready foo",
   897  		Subsystem: structs.NodeEventSubsystemCluster,
   898  		Timestamp: time.Now(),
   899  	}
   900  
   901  	require.NoError(state.UpdateNodeStatus(structs.MsgTypeTestSetup, 801, node.ID, structs.NodeStatusReady, 70, event))
   902  	require.True(watchFired(ws))
   903  
   904  	ws = memdb.NewWatchSet()
   905  	out, err := state.NodeByID(ws, node.ID)
   906  	require.NoError(err)
   907  	require.Equal(structs.NodeStatusReady, out.Status)
   908  	require.EqualValues(801, out.ModifyIndex)
   909  	require.EqualValues(70, out.StatusUpdatedAt)
   910  	require.Len(out.Events, 2)
   911  	require.Equal(event.Message, out.Events[1].Message)
   912  
   913  	index, err := state.Index("nodes")
   914  	require.NoError(err)
   915  	require.EqualValues(801, index)
   916  	require.False(watchFired(ws))
   917  }
   918  
   919  func TestStateStore_BatchUpdateNodeDrain(t *testing.T) {
   920  	t.Parallel()
   921  	require := require.New(t)
   922  
   923  	state := testStateStore(t)
   924  
   925  	n1, n2 := mock.Node(), mock.Node()
   926  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1000, n1))
   927  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1001, n2))
   928  
   929  	// Create a watchset so we can test that update node drain fires the watch
   930  	ws := memdb.NewWatchSet()
   931  	_, err := state.NodeByID(ws, n1.ID)
   932  	require.Nil(err)
   933  
   934  	expectedDrain := &structs.DrainStrategy{
   935  		DrainSpec: structs.DrainSpec{
   936  			Deadline: -1 * time.Second,
   937  		},
   938  	}
   939  
   940  	update := map[string]*structs.DrainUpdate{
   941  		n1.ID: {
   942  			DrainStrategy: expectedDrain,
   943  		},
   944  		n2.ID: {
   945  			DrainStrategy: expectedDrain,
   946  		},
   947  	}
   948  
   949  	event := &structs.NodeEvent{
   950  		Message:   "Drain strategy enabled",
   951  		Subsystem: structs.NodeEventSubsystemDrain,
   952  		Timestamp: time.Now(),
   953  	}
   954  	events := map[string]*structs.NodeEvent{
   955  		n1.ID: event,
   956  		n2.ID: event,
   957  	}
   958  
   959  	require.Nil(state.BatchUpdateNodeDrain(structs.MsgTypeTestSetup, 1002, 7, update, events))
   960  	require.True(watchFired(ws))
   961  
   962  	ws = memdb.NewWatchSet()
   963  	for _, id := range []string{n1.ID, n2.ID} {
   964  		out, err := state.NodeByID(ws, id)
   965  		require.Nil(err)
   966  		require.True(out.Drain)
   967  		require.NotNil(out.DrainStrategy)
   968  		require.Equal(out.DrainStrategy, expectedDrain)
   969  		require.Len(out.Events, 2)
   970  		require.EqualValues(1002, out.ModifyIndex)
   971  		require.EqualValues(7, out.StatusUpdatedAt)
   972  	}
   973  
   974  	index, err := state.Index("nodes")
   975  	require.Nil(err)
   976  	require.EqualValues(1002, index)
   977  	require.False(watchFired(ws))
   978  }
   979  
   980  func TestStateStore_UpdateNodeDrain_Node(t *testing.T) {
   981  	t.Parallel()
   982  	require := require.New(t)
   983  
   984  	state := testStateStore(t)
   985  	node := mock.Node()
   986  
   987  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1000, node))
   988  
   989  	// Create a watchset so we can test that update node drain fires the watch
   990  	ws := memdb.NewWatchSet()
   991  	_, err := state.NodeByID(ws, node.ID)
   992  	require.Nil(err)
   993  
   994  	expectedDrain := &structs.DrainStrategy{
   995  		DrainSpec: structs.DrainSpec{
   996  			Deadline: -1 * time.Second,
   997  		},
   998  	}
   999  
  1000  	event := &structs.NodeEvent{
  1001  		Message:   "Drain strategy enabled",
  1002  		Subsystem: structs.NodeEventSubsystemDrain,
  1003  		Timestamp: time.Now(),
  1004  	}
  1005  	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1001, node.ID, expectedDrain, false, 7, event))
  1006  	require.True(watchFired(ws))
  1007  
  1008  	ws = memdb.NewWatchSet()
  1009  	out, err := state.NodeByID(ws, node.ID)
  1010  	require.Nil(err)
  1011  	require.True(out.Drain)
  1012  	require.NotNil(out.DrainStrategy)
  1013  	require.Equal(out.DrainStrategy, expectedDrain)
  1014  	require.Len(out.Events, 2)
  1015  	require.EqualValues(1001, out.ModifyIndex)
  1016  	require.EqualValues(7, out.StatusUpdatedAt)
  1017  
  1018  	index, err := state.Index("nodes")
  1019  	require.Nil(err)
  1020  	require.EqualValues(1001, index)
  1021  	require.False(watchFired(ws))
  1022  }
  1023  
  1024  func TestStateStore_AddSingleNodeEvent(t *testing.T) {
  1025  	t.Parallel()
  1026  	require := require.New(t)
  1027  
  1028  	state := testStateStore(t)
  1029  
  1030  	node := mock.Node()
  1031  
  1032  	// We create a new node event every time we register a node
  1033  	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
  1034  	require.Nil(err)
  1035  
  1036  	require.Equal(1, len(node.Events))
  1037  	require.Equal(structs.NodeEventSubsystemCluster, node.Events[0].Subsystem)
  1038  	require.Equal(NodeRegisterEventRegistered, node.Events[0].Message)
  1039  
  1040  	// Create a watchset so we can test that AddNodeEvent fires the watch
  1041  	ws := memdb.NewWatchSet()
  1042  	_, err = state.NodeByID(ws, node.ID)
  1043  	require.Nil(err)
  1044  
  1045  	nodeEvent := &structs.NodeEvent{
  1046  		Message:   "failed",
  1047  		Subsystem: "Driver",
  1048  		Timestamp: time.Now(),
  1049  	}
  1050  	nodeEvents := map[string][]*structs.NodeEvent{
  1051  		node.ID: {nodeEvent},
  1052  	}
  1053  	err = state.UpsertNodeEvents(structs.MsgTypeTestSetup, uint64(1001), nodeEvents)
  1054  	require.Nil(err)
  1055  
  1056  	require.True(watchFired(ws))
  1057  
  1058  	ws = memdb.NewWatchSet()
  1059  	out, err := state.NodeByID(ws, node.ID)
  1060  	require.Nil(err)
  1061  
  1062  	require.Equal(2, len(out.Events))
  1063  	require.Equal(nodeEvent, out.Events[1])
  1064  }
  1065  
  1066  // To prevent stale node events from accumulating, we limit the number of
  1067  // stored node events to 10.
  1068  func TestStateStore_NodeEvents_RetentionWindow(t *testing.T) {
  1069  	t.Parallel()
  1070  	require := require.New(t)
  1071  
  1072  	state := testStateStore(t)
  1073  
  1074  	node := mock.Node()
  1075  
  1076  	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
  1077  	if err != nil {
  1078  		t.Fatalf("err: %v", err)
  1079  	}
  1080  	require.Equal(1, len(node.Events))
  1081  	require.Equal(structs.NodeEventSubsystemCluster, node.Events[0].Subsystem)
  1082  	require.Equal(NodeRegisterEventRegistered, node.Events[0].Message)
  1083  
  1084  	var out *structs.Node
  1085  	for i := 1; i <= 20; i++ {
  1086  		ws := memdb.NewWatchSet()
  1087  		out, err = state.NodeByID(ws, node.ID)
  1088  		require.Nil(err)
  1089  
  1090  		nodeEvent := &structs.NodeEvent{
  1091  			Message:   fmt.Sprintf("%dith failed", i),
  1092  			Subsystem: "Driver",
  1093  			Timestamp: time.Now(),
  1094  		}
  1095  
  1096  		nodeEvents := map[string][]*structs.NodeEvent{
  1097  			out.ID: {nodeEvent},
  1098  		}
  1099  		err := state.UpsertNodeEvents(structs.MsgTypeTestSetup, uint64(i), nodeEvents)
  1100  		require.Nil(err)
  1101  
  1102  		require.True(watchFired(ws))
  1103  		ws = memdb.NewWatchSet()
  1104  		out, err = state.NodeByID(ws, node.ID)
  1105  		require.Nil(err)
  1106  	}
  1107  
  1108  	ws := memdb.NewWatchSet()
  1109  	out, err = state.NodeByID(ws, node.ID)
  1110  	require.Nil(err)
  1111  
  1112  	require.Equal(10, len(out.Events))
  1113  	require.Equal(uint64(11), out.Events[0].CreateIndex)
  1114  	require.Equal(uint64(20), out.Events[len(out.Events)-1].CreateIndex)
  1115  }
  1116  
  1117  func TestStateStore_UpdateNodeDrain_ResetEligiblity(t *testing.T) {
  1118  	t.Parallel()
  1119  	require := require.New(t)
  1120  
  1121  	state := testStateStore(t)
  1122  	node := mock.Node()
  1123  	require.Nil(state.UpsertNode(structs.MsgTypeTestSetup, 1000, node))
  1124  
  1125  	// Create a watchset so we can test that update node drain fires the watch
  1126  	ws := memdb.NewWatchSet()
  1127  	_, err := state.NodeByID(ws, node.ID)
  1128  	require.Nil(err)
  1129  
  1130  	drain := &structs.DrainStrategy{
  1131  		DrainSpec: structs.DrainSpec{
  1132  			Deadline: -1 * time.Second,
  1133  		},
  1134  	}
  1135  
  1136  	event1 := &structs.NodeEvent{
  1137  		Message:   "Drain strategy enabled",
  1138  		Subsystem: structs.NodeEventSubsystemDrain,
  1139  		Timestamp: time.Now(),
  1140  	}
  1141  	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1001, node.ID, drain, false, 7, event1))
  1142  	require.True(watchFired(ws))
  1143  
  1144  	// Remove the drain
  1145  	event2 := &structs.NodeEvent{
  1146  		Message:   "Drain strategy disabled",
  1147  		Subsystem: structs.NodeEventSubsystemDrain,
  1148  		Timestamp: time.Now(),
  1149  	}
  1150  	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1002, node.ID, nil, true, 9, event2))
  1151  
  1152  	ws = memdb.NewWatchSet()
  1153  	out, err := state.NodeByID(ws, node.ID)
  1154  	require.Nil(err)
  1155  	require.False(out.Drain)
  1156  	require.Nil(out.DrainStrategy)
  1157  	require.Equal(out.SchedulingEligibility, structs.NodeSchedulingEligible)
  1158  	require.Len(out.Events, 3)
  1159  	require.EqualValues(1002, out.ModifyIndex)
  1160  	require.EqualValues(9, out.StatusUpdatedAt)
  1161  
  1162  	index, err := state.Index("nodes")
  1163  	require.Nil(err)
  1164  	require.EqualValues(1002, index)
  1165  	require.False(watchFired(ws))
  1166  }
  1167  
  1168  func TestStateStore_UpdateNodeEligibility(t *testing.T) {
  1169  	t.Parallel()
  1170  	require := require.New(t)
  1171  
  1172  	state := testStateStore(t)
  1173  	node := mock.Node()
  1174  
  1175  	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
  1176  	if err != nil {
  1177  		t.Fatalf("err: %v", err)
  1178  	}
  1179  
  1180  	expectedEligibility := structs.NodeSchedulingIneligible
  1181  
  1182  	// Create a watchset so we can test that update node drain fires the watch
  1183  	ws := memdb.NewWatchSet()
  1184  	if _, err := state.NodeByID(ws, node.ID); err != nil {
  1185  		t.Fatalf("bad: %v", err)
  1186  	}
  1187  
  1188  	event := &structs.NodeEvent{
  1189  		Message:   "Node marked as ineligible",
  1190  		Subsystem: structs.NodeEventSubsystemCluster,
  1191  		Timestamp: time.Now(),
  1192  	}
  1193  	require.Nil(state.UpdateNodeEligibility(structs.MsgTypeTestSetup, 1001, node.ID, expectedEligibility, 7, event))
  1194  	require.True(watchFired(ws))
  1195  
  1196  	ws = memdb.NewWatchSet()
  1197  	out, err := state.NodeByID(ws, node.ID)
  1198  	require.Nil(err)
  1199  	require.Equal(out.SchedulingEligibility, expectedEligibility)
  1200  	require.Len(out.Events, 2)
  1201  	require.Equal(out.Events[1], event)
  1202  	require.EqualValues(1001, out.ModifyIndex)
  1203  	require.EqualValues(7, out.StatusUpdatedAt)
  1204  
  1205  	index, err := state.Index("nodes")
  1206  	require.Nil(err)
  1207  	require.EqualValues(1001, index)
  1208  	require.False(watchFired(ws))
  1209  
  1210  	// Set a drain strategy
  1211  	expectedDrain := &structs.DrainStrategy{
  1212  		DrainSpec: structs.DrainSpec{
  1213  			Deadline: -1 * time.Second,
  1214  		},
  1215  	}
  1216  	require.Nil(state.UpdateNodeDrain(structs.MsgTypeTestSetup, 1002, node.ID, expectedDrain, false, 7, nil))
  1217  
  1218  	// Try to set the node to eligible
  1219  	err = state.UpdateNodeEligibility(structs.MsgTypeTestSetup, 1003, node.ID, structs.NodeSchedulingEligible, 9, nil)
  1220  	require.NotNil(err)
  1221  	require.Contains(err.Error(), "while it is draining")
  1222  }
  1223  
  1224  func TestStateStore_Nodes(t *testing.T) {
  1225  	t.Parallel()
  1226  
  1227  	state := testStateStore(t)
  1228  	var nodes []*structs.Node
  1229  
  1230  	for i := 0; i < 10; i++ {
  1231  		node := mock.Node()
  1232  		nodes = append(nodes, node)
  1233  
  1234  		err := state.UpsertNode(structs.MsgTypeTestSetup, 1000+uint64(i), node)
  1235  		if err != nil {
  1236  			t.Fatalf("err: %v", err)
  1237  		}
  1238  	}
  1239  
  1240  	// Create a watchset so we can test that getters don't cause it to fire
  1241  	ws := memdb.NewWatchSet()
  1242  	iter, err := state.Nodes(ws)
  1243  	if err != nil {
  1244  		t.Fatalf("bad: %v", err)
  1245  	}
  1246  
  1247  	var out []*structs.Node
  1248  	for {
  1249  		raw := iter.Next()
  1250  		if raw == nil {
  1251  			break
  1252  		}
  1253  		out = append(out, raw.(*structs.Node))
  1254  	}
  1255  
  1256  	sort.Sort(NodeIDSort(nodes))
  1257  	sort.Sort(NodeIDSort(out))
  1258  
  1259  	if !reflect.DeepEqual(nodes, out) {
  1260  		t.Fatalf("bad: %#v %#v", nodes, out)
  1261  	}
  1262  
  1263  	if watchFired(ws) {
  1264  		t.Fatalf("bad")
  1265  	}
  1266  }
  1267  
  1268  func TestStateStore_NodesByIDPrefix(t *testing.T) {
  1269  	t.Parallel()
  1270  
  1271  	state := testStateStore(t)
  1272  	node := mock.Node()
  1273  
  1274  	node.ID = "11111111-662e-d0ab-d1c9-3e434af7bdb4"
  1275  	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
  1276  	if err != nil {
  1277  		t.Fatalf("err: %v", err)
  1278  	}
  1279  
  1280  	// Create a watchset so we can test that getters don't cause it to fire
  1281  	ws := memdb.NewWatchSet()
  1282  	iter, err := state.NodesByIDPrefix(ws, node.ID)
  1283  	if err != nil {
  1284  		t.Fatalf("err: %v", err)
  1285  	}
  1286  
  1287  	gatherNodes := func(iter memdb.ResultIterator) []*structs.Node {
  1288  		var nodes []*structs.Node
  1289  		for {
  1290  			raw := iter.Next()
  1291  			if raw == nil {
  1292  				break
  1293  			}
  1294  			node := raw.(*structs.Node)
  1295  			nodes = append(nodes, node)
  1296  		}
  1297  		return nodes
  1298  	}
  1299  
  1300  	nodes := gatherNodes(iter)
  1301  	if len(nodes) != 1 {
  1302  		t.Fatalf("err: %v", err)
  1303  	}
  1304  
  1305  	if watchFired(ws) {
  1306  		t.Fatalf("bad")
  1307  	}
  1308  
  1309  	iter, err = state.NodesByIDPrefix(ws, "11")
  1310  	if err != nil {
  1311  		t.Fatalf("err: %v", err)
  1312  	}
  1313  
  1314  	nodes = gatherNodes(iter)
  1315  	if len(nodes) != 1 {
  1316  		t.Fatalf("err: %v", err)
  1317  	}
  1318  
  1319  	node = mock.Node()
  1320  	node.ID = "11222222-662e-d0ab-d1c9-3e434af7bdb4"
  1321  	err = state.UpsertNode(structs.MsgTypeTestSetup, 1001, node)
  1322  	if err != nil {
  1323  		t.Fatalf("err: %v", err)
  1324  	}
  1325  
  1326  	if !watchFired(ws) {
  1327  		t.Fatalf("bad")
  1328  	}
  1329  
  1330  	ws = memdb.NewWatchSet()
  1331  	iter, err = state.NodesByIDPrefix(ws, "11")
  1332  	if err != nil {
  1333  		t.Fatalf("err: %v", err)
  1334  	}
  1335  
  1336  	nodes = gatherNodes(iter)
  1337  	if len(nodes) != 2 {
  1338  		t.Fatalf("err: %v", err)
  1339  	}
  1340  
  1341  	iter, err = state.NodesByIDPrefix(ws, "1111")
  1342  	if err != nil {
  1343  		t.Fatalf("err: %v", err)
  1344  	}
  1345  
  1346  	nodes = gatherNodes(iter)
  1347  	if len(nodes) != 1 {
  1348  		t.Fatalf("err: %v", err)
  1349  	}
  1350  
  1351  	if watchFired(ws) {
  1352  		t.Fatalf("bad")
  1353  	}
  1354  }
  1355  
  1356  func TestStateStore_RestoreNode(t *testing.T) {
  1357  	t.Parallel()
  1358  
  1359  	state := testStateStore(t)
  1360  	node := mock.Node()
  1361  
  1362  	restore, err := state.Restore()
  1363  	if err != nil {
  1364  		t.Fatalf("err: %v", err)
  1365  	}
  1366  
  1367  	err = restore.NodeRestore(node)
  1368  	if err != nil {
  1369  		t.Fatalf("err: %v", err)
  1370  	}
  1371  	require.NoError(t, restore.Commit())
  1372  
  1373  	ws := memdb.NewWatchSet()
  1374  	out, err := state.NodeByID(ws, node.ID)
  1375  	if err != nil {
  1376  		t.Fatalf("err: %v", err)
  1377  	}
  1378  
  1379  	if !reflect.DeepEqual(out, node) {
  1380  		t.Fatalf("Bad: %#v %#v", out, node)
  1381  	}
  1382  }
  1383  
  1384  func TestStateStore_UpsertJob_Job(t *testing.T) {
  1385  	t.Parallel()
  1386  
  1387  	state := testStateStore(t)
  1388  	job := mock.Job()
  1389  
  1390  	// Create a watchset so we can test that upsert fires the watch
  1391  	ws := memdb.NewWatchSet()
  1392  	_, err := state.JobByID(ws, job.Namespace, job.ID)
  1393  	if err != nil {
  1394  		t.Fatalf("bad: %v", err)
  1395  	}
  1396  
  1397  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
  1398  		t.Fatalf("err: %v", err)
  1399  	}
  1400  	if !watchFired(ws) {
  1401  		t.Fatalf("bad")
  1402  	}
  1403  
  1404  	ws = memdb.NewWatchSet()
  1405  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1406  	if err != nil {
  1407  		t.Fatalf("err: %v", err)
  1408  	}
  1409  
  1410  	if !reflect.DeepEqual(job, out) {
  1411  		t.Fatalf("bad: %#v %#v", job, out)
  1412  	}
  1413  
  1414  	index, err := state.Index("jobs")
  1415  	if err != nil {
  1416  		t.Fatalf("err: %v", err)
  1417  	}
  1418  	if index != 1000 {
  1419  		t.Fatalf("bad: %d", index)
  1420  	}
  1421  
  1422  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1423  	if err != nil {
  1424  		t.Fatalf("err: %v", err)
  1425  	}
  1426  	if summary == nil {
  1427  		t.Fatalf("nil summary")
  1428  	}
  1429  	if summary.JobID != job.ID {
  1430  		t.Fatalf("bad summary id: %v", summary.JobID)
  1431  	}
  1432  	_, ok := summary.Summary["web"]
  1433  	if !ok {
  1434  		t.Fatalf("nil summary for task group")
  1435  	}
  1436  	if watchFired(ws) {
  1437  		t.Fatalf("bad")
  1438  	}
  1439  
  1440  	// Check the job versions
  1441  	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1442  	if err != nil {
  1443  		t.Fatalf("err: %v", err)
  1444  	}
  1445  	if len(allVersions) != 1 {
  1446  		t.Fatalf("got %d; want 1", len(allVersions))
  1447  	}
  1448  
  1449  	if a := allVersions[0]; a.ID != job.ID || a.Version != 0 {
  1450  		t.Fatalf("bad: %v", a)
  1451  	}
  1452  
  1453  	// Test the looking up the job by version returns the same results
  1454  	vout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
  1455  	if err != nil {
  1456  		t.Fatalf("err: %v", err)
  1457  	}
  1458  
  1459  	if !reflect.DeepEqual(out, vout) {
  1460  		t.Fatalf("bad: %#v %#v", out, vout)
  1461  	}
  1462  }
  1463  
  1464  func TestStateStore_UpdateUpsertJob_Job(t *testing.T) {
  1465  	t.Parallel()
  1466  
  1467  	state := testStateStore(t)
  1468  	job := mock.Job()
  1469  
  1470  	// Create a watchset so we can test that upsert fires the watch
  1471  	ws := memdb.NewWatchSet()
  1472  	_, err := state.JobByID(ws, job.Namespace, job.ID)
  1473  	if err != nil {
  1474  		t.Fatalf("bad: %v", err)
  1475  	}
  1476  
  1477  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
  1478  		t.Fatalf("err: %v", err)
  1479  	}
  1480  
  1481  	job2 := mock.Job()
  1482  	job2.ID = job.ID
  1483  	job2.AllAtOnce = true
  1484  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1001, job2)
  1485  	if err != nil {
  1486  		t.Fatalf("err: %v", err)
  1487  	}
  1488  
  1489  	if !watchFired(ws) {
  1490  		t.Fatalf("bad")
  1491  	}
  1492  
  1493  	ws = memdb.NewWatchSet()
  1494  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1495  	if err != nil {
  1496  		t.Fatalf("err: %v", err)
  1497  	}
  1498  
  1499  	if !reflect.DeepEqual(job2, out) {
  1500  		t.Fatalf("bad: %#v %#v", job2, out)
  1501  	}
  1502  
  1503  	if out.CreateIndex != 1000 {
  1504  		t.Fatalf("bad: %#v", out)
  1505  	}
  1506  	if out.ModifyIndex != 1001 {
  1507  		t.Fatalf("bad: %#v", out)
  1508  	}
  1509  	if out.Version != 1 {
  1510  		t.Fatalf("bad: %#v", out)
  1511  	}
  1512  
  1513  	index, err := state.Index("jobs")
  1514  	if err != nil {
  1515  		t.Fatalf("err: %v", err)
  1516  	}
  1517  	if index != 1001 {
  1518  		t.Fatalf("bad: %d", index)
  1519  	}
  1520  
  1521  	// Test the looking up the job by version returns the same results
  1522  	vout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 1)
  1523  	if err != nil {
  1524  		t.Fatalf("err: %v", err)
  1525  	}
  1526  
  1527  	if !reflect.DeepEqual(out, vout) {
  1528  		t.Fatalf("bad: %#v %#v", out, vout)
  1529  	}
  1530  
  1531  	// Test that the job summary remains the same if the job is updated but
  1532  	// count remains same
  1533  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1534  	if err != nil {
  1535  		t.Fatalf("err: %v", err)
  1536  	}
  1537  	if summary == nil {
  1538  		t.Fatalf("nil summary")
  1539  	}
  1540  	if summary.JobID != job.ID {
  1541  		t.Fatalf("bad summary id: %v", summary.JobID)
  1542  	}
  1543  	_, ok := summary.Summary["web"]
  1544  	if !ok {
  1545  		t.Fatalf("nil summary for task group")
  1546  	}
  1547  
  1548  	// Check the job versions
  1549  	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1550  	if err != nil {
  1551  		t.Fatalf("err: %v", err)
  1552  	}
  1553  	if len(allVersions) != 2 {
  1554  		t.Fatalf("got %d; want 1", len(allVersions))
  1555  	}
  1556  
  1557  	if a := allVersions[0]; a.ID != job.ID || a.Version != 1 || !a.AllAtOnce {
  1558  		t.Fatalf("bad: %+v", a)
  1559  	}
  1560  	if a := allVersions[1]; a.ID != job.ID || a.Version != 0 || a.AllAtOnce {
  1561  		t.Fatalf("bad: %+v", a)
  1562  	}
  1563  
  1564  	if watchFired(ws) {
  1565  		t.Fatalf("bad")
  1566  	}
  1567  }
  1568  
  1569  func TestStateStore_UpdateUpsertJob_PeriodicJob(t *testing.T) {
  1570  	t.Parallel()
  1571  
  1572  	state := testStateStore(t)
  1573  	job := mock.PeriodicJob()
  1574  
  1575  	// Create a watchset so we can test that upsert fires the watch
  1576  	ws := memdb.NewWatchSet()
  1577  	_, err := state.JobByID(ws, job.Namespace, job.ID)
  1578  	if err != nil {
  1579  		t.Fatalf("bad: %v", err)
  1580  	}
  1581  
  1582  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
  1583  		t.Fatalf("err: %v", err)
  1584  	}
  1585  
  1586  	// Create a child and an evaluation
  1587  	job2 := job.Copy()
  1588  	job2.Periodic = nil
  1589  	job2.ID = fmt.Sprintf("%v/%s-1490635020", job.ID, structs.PeriodicLaunchSuffix)
  1590  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1001, job2)
  1591  	if err != nil {
  1592  		t.Fatalf("err: %v", err)
  1593  	}
  1594  
  1595  	eval := mock.Eval()
  1596  	eval.JobID = job2.ID
  1597  	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1002, []*structs.Evaluation{eval})
  1598  	if err != nil {
  1599  		t.Fatalf("err: %v", err)
  1600  	}
  1601  
  1602  	job3 := job.Copy()
  1603  	job3.TaskGroups[0].Tasks[0].Name = "new name"
  1604  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1003, job3)
  1605  	if err != nil {
  1606  		t.Fatalf("err: %v", err)
  1607  	}
  1608  
  1609  	if !watchFired(ws) {
  1610  		t.Fatalf("bad")
  1611  	}
  1612  
  1613  	ws = memdb.NewWatchSet()
  1614  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1615  	if err != nil {
  1616  		t.Fatalf("err: %v", err)
  1617  	}
  1618  
  1619  	if s, e := out.Status, structs.JobStatusRunning; s != e {
  1620  		t.Fatalf("got status %v; want %v", s, e)
  1621  	}
  1622  
  1623  }
  1624  
  1625  func TestStateStore_UpsertJob_BadNamespace(t *testing.T) {
  1626  	t.Parallel()
  1627  
  1628  	assert := assert.New(t)
  1629  	state := testStateStore(t)
  1630  	job := mock.Job()
  1631  	job.Namespace = "foo"
  1632  
  1633  	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  1634  	assert.Contains(err.Error(), "nonexistent namespace")
  1635  
  1636  	ws := memdb.NewWatchSet()
  1637  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1638  	assert.Nil(err)
  1639  	assert.Nil(out)
  1640  }
  1641  
  1642  // Upsert a job that is the child of a parent job and ensures its summary gets
  1643  // updated.
  1644  func TestStateStore_UpsertJob_ChildJob(t *testing.T) {
  1645  	t.Parallel()
  1646  
  1647  	state := testStateStore(t)
  1648  
  1649  	// Create a watchset so we can test that upsert fires the watch
  1650  	parent := mock.Job()
  1651  	ws := memdb.NewWatchSet()
  1652  	_, err := state.JobByID(ws, parent.Namespace, parent.ID)
  1653  	if err != nil {
  1654  		t.Fatalf("bad: %v", err)
  1655  	}
  1656  
  1657  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, parent); err != nil {
  1658  		t.Fatalf("err: %v", err)
  1659  	}
  1660  
  1661  	child := mock.Job()
  1662  	child.Status = ""
  1663  	child.ParentID = parent.ID
  1664  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1001, child); err != nil {
  1665  		t.Fatalf("err: %v", err)
  1666  	}
  1667  
  1668  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  1669  	if err != nil {
  1670  		t.Fatalf("err: %v", err)
  1671  	}
  1672  	if summary == nil {
  1673  		t.Fatalf("nil summary")
  1674  	}
  1675  	if summary.JobID != parent.ID {
  1676  		t.Fatalf("bad summary id: %v", parent.ID)
  1677  	}
  1678  	if summary.Children == nil {
  1679  		t.Fatalf("nil children summary")
  1680  	}
  1681  	if summary.Children.Pending != 1 || summary.Children.Running != 0 || summary.Children.Dead != 0 {
  1682  		t.Fatalf("bad children summary: %v", summary.Children)
  1683  	}
  1684  	if !watchFired(ws) {
  1685  		t.Fatalf("bad")
  1686  	}
  1687  }
  1688  
  1689  func TestStateStore_UpdateUpsertJob_JobVersion(t *testing.T) {
  1690  	t.Parallel()
  1691  
  1692  	state := testStateStore(t)
  1693  
  1694  	// Create a job and mark it as stable
  1695  	job := mock.Job()
  1696  	job.Stable = true
  1697  	job.Name = "0"
  1698  
  1699  	// Create a watchset so we can test that upsert fires the watch
  1700  	ws := memdb.NewWatchSet()
  1701  	_, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1702  	if err != nil {
  1703  		t.Fatalf("bad: %v", err)
  1704  	}
  1705  
  1706  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
  1707  		t.Fatalf("err: %v", err)
  1708  	}
  1709  
  1710  	if !watchFired(ws) {
  1711  		t.Fatalf("bad")
  1712  	}
  1713  
  1714  	var finalJob *structs.Job
  1715  	for i := 1; i < 300; i++ {
  1716  		finalJob = mock.Job()
  1717  		finalJob.ID = job.ID
  1718  		finalJob.Name = fmt.Sprintf("%d", i)
  1719  		err = state.UpsertJob(structs.MsgTypeTestSetup, uint64(1000+i), finalJob)
  1720  		if err != nil {
  1721  			t.Fatalf("err: %v", err)
  1722  		}
  1723  	}
  1724  
  1725  	ws = memdb.NewWatchSet()
  1726  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1727  	if err != nil {
  1728  		t.Fatalf("err: %v", err)
  1729  	}
  1730  
  1731  	if !reflect.DeepEqual(finalJob, out) {
  1732  		t.Fatalf("bad: %#v %#v", finalJob, out)
  1733  	}
  1734  
  1735  	if out.CreateIndex != 1000 {
  1736  		t.Fatalf("bad: %#v", out)
  1737  	}
  1738  	if out.ModifyIndex != 1299 {
  1739  		t.Fatalf("bad: %#v", out)
  1740  	}
  1741  	if out.Version != 299 {
  1742  		t.Fatalf("bad: %#v", out)
  1743  	}
  1744  
  1745  	index, err := state.Index("job_version")
  1746  	if err != nil {
  1747  		t.Fatalf("err: %v", err)
  1748  	}
  1749  	if index != 1299 {
  1750  		t.Fatalf("bad: %d", index)
  1751  	}
  1752  
  1753  	// Check the job versions
  1754  	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1755  	if err != nil {
  1756  		t.Fatalf("err: %v", err)
  1757  	}
  1758  	if len(allVersions) != structs.JobTrackedVersions {
  1759  		t.Fatalf("got %d; want %d", len(allVersions), structs.JobTrackedVersions)
  1760  	}
  1761  
  1762  	if a := allVersions[0]; a.ID != job.ID || a.Version != 299 || a.Name != "299" {
  1763  		t.Fatalf("bad: %+v", a)
  1764  	}
  1765  	if a := allVersions[1]; a.ID != job.ID || a.Version != 298 || a.Name != "298" {
  1766  		t.Fatalf("bad: %+v", a)
  1767  	}
  1768  
  1769  	// Ensure we didn't delete the stable job
  1770  	if a := allVersions[structs.JobTrackedVersions-1]; a.ID != job.ID ||
  1771  		a.Version != 0 || a.Name != "0" || !a.Stable {
  1772  		t.Fatalf("bad: %+v", a)
  1773  	}
  1774  
  1775  	if watchFired(ws) {
  1776  		t.Fatalf("bad")
  1777  	}
  1778  }
  1779  
  1780  func TestStateStore_DeleteJob_Job(t *testing.T) {
  1781  	t.Parallel()
  1782  
  1783  	state := testStateStore(t)
  1784  	job := mock.Job()
  1785  
  1786  	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  1787  	if err != nil {
  1788  		t.Fatalf("err: %v", err)
  1789  	}
  1790  
  1791  	// Create a watchset so we can test that delete fires the watch
  1792  	ws := memdb.NewWatchSet()
  1793  	if _, err := state.JobByID(ws, job.Namespace, job.ID); err != nil {
  1794  		t.Fatalf("bad: %v", err)
  1795  	}
  1796  
  1797  	err = state.DeleteJob(1001, job.Namespace, job.ID)
  1798  	if err != nil {
  1799  		t.Fatalf("err: %v", err)
  1800  	}
  1801  
  1802  	if !watchFired(ws) {
  1803  		t.Fatalf("bad")
  1804  	}
  1805  
  1806  	ws = memdb.NewWatchSet()
  1807  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1808  	if err != nil {
  1809  		t.Fatalf("err: %v", err)
  1810  	}
  1811  
  1812  	if out != nil {
  1813  		t.Fatalf("bad: %#v %#v", job, out)
  1814  	}
  1815  
  1816  	index, err := state.Index("jobs")
  1817  	if err != nil {
  1818  		t.Fatalf("err: %v", err)
  1819  	}
  1820  	if index != 1001 {
  1821  		t.Fatalf("bad: %d", index)
  1822  	}
  1823  
  1824  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1825  	if err != nil {
  1826  		t.Fatalf("err: %v", err)
  1827  	}
  1828  	if summary != nil {
  1829  		t.Fatalf("expected summary to be nil, but got: %v", summary)
  1830  	}
  1831  
  1832  	index, err = state.Index("job_summary")
  1833  	if err != nil {
  1834  		t.Fatalf("err: %v", err)
  1835  	}
  1836  	if index != 1001 {
  1837  		t.Fatalf("bad: %d", index)
  1838  	}
  1839  
  1840  	versions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1841  	if err != nil {
  1842  		t.Fatalf("err: %v", err)
  1843  	}
  1844  	if len(versions) != 0 {
  1845  		t.Fatalf("expected no job versions")
  1846  	}
  1847  
  1848  	index, err = state.Index("job_summary")
  1849  	if err != nil {
  1850  		t.Fatalf("err: %v", err)
  1851  	}
  1852  	if index != 1001 {
  1853  		t.Fatalf("bad: %d", index)
  1854  	}
  1855  
  1856  	if watchFired(ws) {
  1857  		t.Fatalf("bad")
  1858  	}
  1859  }
  1860  
  1861  func TestStateStore_DeleteJobTxn_BatchDeletes(t *testing.T) {
  1862  	t.Parallel()
  1863  
  1864  	state := testStateStore(t)
  1865  
  1866  	const testJobCount = 10
  1867  	const jobVersionCount = 4
  1868  
  1869  	stateIndex := uint64(1000)
  1870  
  1871  	jobs := make([]*structs.Job, testJobCount)
  1872  	for i := 0; i < testJobCount; i++ {
  1873  		stateIndex++
  1874  		job := mock.BatchJob()
  1875  
  1876  		err := state.UpsertJob(structs.MsgTypeTestSetup, stateIndex, job)
  1877  		require.NoError(t, err)
  1878  
  1879  		jobs[i] = job
  1880  
  1881  		// Create some versions
  1882  		for vi := 1; vi < jobVersionCount; vi++ {
  1883  			stateIndex++
  1884  
  1885  			job := job.Copy()
  1886  			job.TaskGroups[0].Tasks[0].Env = map[string]string{
  1887  				"Version": fmt.Sprintf("%d", vi),
  1888  			}
  1889  
  1890  			require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, stateIndex, job))
  1891  		}
  1892  	}
  1893  
  1894  	ws := memdb.NewWatchSet()
  1895  
  1896  	// Sanity check that jobs are present in DB
  1897  	job, err := state.JobByID(ws, jobs[0].Namespace, jobs[0].ID)
  1898  	require.NoError(t, err)
  1899  	require.Equal(t, jobs[0].ID, job.ID)
  1900  
  1901  	jobVersions, err := state.JobVersionsByID(ws, jobs[0].Namespace, jobs[0].ID)
  1902  	require.NoError(t, err)
  1903  	require.Equal(t, jobVersionCount, len(jobVersions))
  1904  
  1905  	// Actually delete
  1906  	const deletionIndex = uint64(10001)
  1907  	err = state.WithWriteTransaction(structs.MsgTypeTestSetup, deletionIndex, func(txn Txn) error {
  1908  		for i, job := range jobs {
  1909  			err := state.DeleteJobTxn(deletionIndex, job.Namespace, job.ID, txn)
  1910  			require.NoError(t, err, "failed at %d %e", i, err)
  1911  		}
  1912  		return nil
  1913  	})
  1914  	assert.NoError(t, err)
  1915  
  1916  	assert.True(t, watchFired(ws))
  1917  
  1918  	ws = memdb.NewWatchSet()
  1919  	out, err := state.JobByID(ws, jobs[0].Namespace, jobs[0].ID)
  1920  	require.NoError(t, err)
  1921  	require.Nil(t, out)
  1922  
  1923  	jobVersions, err = state.JobVersionsByID(ws, jobs[0].Namespace, jobs[0].ID)
  1924  	require.NoError(t, err)
  1925  	require.Empty(t, jobVersions)
  1926  
  1927  	index, err := state.Index("jobs")
  1928  	require.NoError(t, err)
  1929  	require.Equal(t, deletionIndex, index)
  1930  }
  1931  
  1932  func TestStateStore_DeleteJob_MultipleVersions(t *testing.T) {
  1933  	t.Parallel()
  1934  
  1935  	state := testStateStore(t)
  1936  	assert := assert.New(t)
  1937  
  1938  	// Create a job and mark it as stable
  1939  	job := mock.Job()
  1940  	job.Stable = true
  1941  	job.Priority = 0
  1942  
  1943  	// Create a watchset so we can test that upsert fires the watch
  1944  	ws := memdb.NewWatchSet()
  1945  	_, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1946  	assert.Nil(err)
  1947  	assert.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1000, job))
  1948  	assert.True(watchFired(ws))
  1949  
  1950  	var finalJob *structs.Job
  1951  	for i := 1; i < 20; i++ {
  1952  		finalJob = mock.Job()
  1953  		finalJob.ID = job.ID
  1954  		finalJob.Priority = i
  1955  		assert.Nil(state.UpsertJob(structs.MsgTypeTestSetup, uint64(1000+i), finalJob))
  1956  	}
  1957  
  1958  	assert.Nil(state.DeleteJob(1020, job.Namespace, job.ID))
  1959  	assert.True(watchFired(ws))
  1960  
  1961  	ws = memdb.NewWatchSet()
  1962  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1963  	assert.Nil(err)
  1964  	assert.Nil(out)
  1965  
  1966  	index, err := state.Index("jobs")
  1967  	assert.Nil(err)
  1968  	assert.EqualValues(1020, index)
  1969  
  1970  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1971  	assert.Nil(err)
  1972  	assert.Nil(summary)
  1973  
  1974  	index, err = state.Index("job_version")
  1975  	assert.Nil(err)
  1976  	assert.EqualValues(1020, index)
  1977  
  1978  	versions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1979  	assert.Nil(err)
  1980  	assert.Len(versions, 0)
  1981  
  1982  	index, err = state.Index("job_summary")
  1983  	assert.Nil(err)
  1984  	assert.EqualValues(1020, index)
  1985  
  1986  	assert.False(watchFired(ws))
  1987  }
  1988  
  1989  func TestStateStore_DeleteJob_ChildJob(t *testing.T) {
  1990  	t.Parallel()
  1991  
  1992  	state := testStateStore(t)
  1993  
  1994  	parent := mock.Job()
  1995  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
  1996  		t.Fatalf("err: %v", err)
  1997  	}
  1998  
  1999  	child := mock.Job()
  2000  	child.Status = ""
  2001  	child.ParentID = parent.ID
  2002  
  2003  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
  2004  		t.Fatalf("err: %v", err)
  2005  	}
  2006  
  2007  	// Create a watchset so we can test that delete fires the watch
  2008  	ws := memdb.NewWatchSet()
  2009  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  2010  		t.Fatalf("bad: %v", err)
  2011  	}
  2012  
  2013  	err := state.DeleteJob(1001, child.Namespace, child.ID)
  2014  	if err != nil {
  2015  		t.Fatalf("err: %v", err)
  2016  	}
  2017  	if !watchFired(ws) {
  2018  		t.Fatalf("bad")
  2019  	}
  2020  
  2021  	ws = memdb.NewWatchSet()
  2022  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  2023  	if err != nil {
  2024  		t.Fatalf("err: %v", err)
  2025  	}
  2026  	if summary == nil {
  2027  		t.Fatalf("nil summary")
  2028  	}
  2029  	if summary.JobID != parent.ID {
  2030  		t.Fatalf("bad summary id: %v", parent.ID)
  2031  	}
  2032  	if summary.Children == nil {
  2033  		t.Fatalf("nil children summary")
  2034  	}
  2035  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  2036  		t.Fatalf("bad children summary: %v", summary.Children)
  2037  	}
  2038  	if watchFired(ws) {
  2039  		t.Fatalf("bad")
  2040  	}
  2041  }
  2042  
  2043  func TestStateStore_Jobs(t *testing.T) {
  2044  	t.Parallel()
  2045  
  2046  	state := testStateStore(t)
  2047  	var jobs []*structs.Job
  2048  
  2049  	for i := 0; i < 10; i++ {
  2050  		job := mock.Job()
  2051  		jobs = append(jobs, job)
  2052  
  2053  		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
  2054  		if err != nil {
  2055  			t.Fatalf("err: %v", err)
  2056  		}
  2057  	}
  2058  
  2059  	ws := memdb.NewWatchSet()
  2060  	iter, err := state.Jobs(ws)
  2061  	if err != nil {
  2062  		t.Fatalf("err: %v", err)
  2063  	}
  2064  
  2065  	var out []*structs.Job
  2066  	for {
  2067  		raw := iter.Next()
  2068  		if raw == nil {
  2069  			break
  2070  		}
  2071  		out = append(out, raw.(*structs.Job))
  2072  	}
  2073  
  2074  	sort.Sort(JobIDSort(jobs))
  2075  	sort.Sort(JobIDSort(out))
  2076  
  2077  	if !reflect.DeepEqual(jobs, out) {
  2078  		t.Fatalf("bad: %#v %#v", jobs, out)
  2079  	}
  2080  	if watchFired(ws) {
  2081  		t.Fatalf("bad")
  2082  	}
  2083  }
  2084  
  2085  func TestStateStore_JobVersions(t *testing.T) {
  2086  	t.Parallel()
  2087  
  2088  	state := testStateStore(t)
  2089  	var jobs []*structs.Job
  2090  
  2091  	for i := 0; i < 10; i++ {
  2092  		job := mock.Job()
  2093  		jobs = append(jobs, job)
  2094  
  2095  		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
  2096  		if err != nil {
  2097  			t.Fatalf("err: %v", err)
  2098  		}
  2099  	}
  2100  
  2101  	ws := memdb.NewWatchSet()
  2102  	iter, err := state.JobVersions(ws)
  2103  	if err != nil {
  2104  		t.Fatalf("err: %v", err)
  2105  	}
  2106  
  2107  	var out []*structs.Job
  2108  	for {
  2109  		raw := iter.Next()
  2110  		if raw == nil {
  2111  			break
  2112  		}
  2113  		out = append(out, raw.(*structs.Job))
  2114  	}
  2115  
  2116  	sort.Sort(JobIDSort(jobs))
  2117  	sort.Sort(JobIDSort(out))
  2118  
  2119  	if !reflect.DeepEqual(jobs, out) {
  2120  		t.Fatalf("bad: %#v %#v", jobs, out)
  2121  	}
  2122  	if watchFired(ws) {
  2123  		t.Fatalf("bad")
  2124  	}
  2125  }
  2126  
  2127  func TestStateStore_JobsByIDPrefix(t *testing.T) {
  2128  	t.Parallel()
  2129  
  2130  	state := testStateStore(t)
  2131  	job := mock.Job()
  2132  
  2133  	job.ID = "redis"
  2134  	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  2135  	if err != nil {
  2136  		t.Fatalf("err: %v", err)
  2137  	}
  2138  
  2139  	ws := memdb.NewWatchSet()
  2140  	iter, err := state.JobsByIDPrefix(ws, job.Namespace, job.ID)
  2141  	if err != nil {
  2142  		t.Fatalf("err: %v", err)
  2143  	}
  2144  
  2145  	gatherJobs := func(iter memdb.ResultIterator) []*structs.Job {
  2146  		var jobs []*structs.Job
  2147  		for {
  2148  			raw := iter.Next()
  2149  			if raw == nil {
  2150  				break
  2151  			}
  2152  			jobs = append(jobs, raw.(*structs.Job))
  2153  		}
  2154  		return jobs
  2155  	}
  2156  
  2157  	jobs := gatherJobs(iter)
  2158  	if len(jobs) != 1 {
  2159  		t.Fatalf("err: %v", err)
  2160  	}
  2161  
  2162  	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "re")
  2163  	if err != nil {
  2164  		t.Fatalf("err: %v", err)
  2165  	}
  2166  
  2167  	jobs = gatherJobs(iter)
  2168  	if len(jobs) != 1 {
  2169  		t.Fatalf("err: %v", err)
  2170  	}
  2171  	if watchFired(ws) {
  2172  		t.Fatalf("bad")
  2173  	}
  2174  
  2175  	job = mock.Job()
  2176  	job.ID = "riak"
  2177  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1001, job)
  2178  	if err != nil {
  2179  		t.Fatalf("err: %v", err)
  2180  	}
  2181  
  2182  	if !watchFired(ws) {
  2183  		t.Fatalf("bad")
  2184  	}
  2185  
  2186  	ws = memdb.NewWatchSet()
  2187  	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "r")
  2188  	if err != nil {
  2189  		t.Fatalf("err: %v", err)
  2190  	}
  2191  
  2192  	jobs = gatherJobs(iter)
  2193  	if len(jobs) != 2 {
  2194  		t.Fatalf("err: %v", err)
  2195  	}
  2196  
  2197  	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "ri")
  2198  	if err != nil {
  2199  		t.Fatalf("err: %v", err)
  2200  	}
  2201  
  2202  	jobs = gatherJobs(iter)
  2203  	if len(jobs) != 1 {
  2204  		t.Fatalf("err: %v", err)
  2205  	}
  2206  	if watchFired(ws) {
  2207  		t.Fatalf("bad")
  2208  	}
  2209  }
  2210  
  2211  func TestStateStore_JobsByPeriodic(t *testing.T) {
  2212  	t.Parallel()
  2213  
  2214  	state := testStateStore(t)
  2215  	var periodic, nonPeriodic []*structs.Job
  2216  
  2217  	for i := 0; i < 10; i++ {
  2218  		job := mock.Job()
  2219  		nonPeriodic = append(nonPeriodic, job)
  2220  
  2221  		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
  2222  		if err != nil {
  2223  			t.Fatalf("err: %v", err)
  2224  		}
  2225  	}
  2226  
  2227  	for i := 0; i < 10; i++ {
  2228  		job := mock.PeriodicJob()
  2229  		periodic = append(periodic, job)
  2230  
  2231  		err := state.UpsertJob(structs.MsgTypeTestSetup, 2000+uint64(i), job)
  2232  		if err != nil {
  2233  			t.Fatalf("err: %v", err)
  2234  		}
  2235  	}
  2236  
  2237  	ws := memdb.NewWatchSet()
  2238  	iter, err := state.JobsByPeriodic(ws, true)
  2239  	if err != nil {
  2240  		t.Fatalf("err: %v", err)
  2241  	}
  2242  
  2243  	var outPeriodic []*structs.Job
  2244  	for {
  2245  		raw := iter.Next()
  2246  		if raw == nil {
  2247  			break
  2248  		}
  2249  		outPeriodic = append(outPeriodic, raw.(*structs.Job))
  2250  	}
  2251  
  2252  	iter, err = state.JobsByPeriodic(ws, false)
  2253  	if err != nil {
  2254  		t.Fatalf("err: %v", err)
  2255  	}
  2256  
  2257  	var outNonPeriodic []*structs.Job
  2258  	for {
  2259  		raw := iter.Next()
  2260  		if raw == nil {
  2261  			break
  2262  		}
  2263  		outNonPeriodic = append(outNonPeriodic, raw.(*structs.Job))
  2264  	}
  2265  
  2266  	sort.Sort(JobIDSort(periodic))
  2267  	sort.Sort(JobIDSort(nonPeriodic))
  2268  	sort.Sort(JobIDSort(outPeriodic))
  2269  	sort.Sort(JobIDSort(outNonPeriodic))
  2270  
  2271  	if !reflect.DeepEqual(periodic, outPeriodic) {
  2272  		t.Fatalf("bad: %#v %#v", periodic, outPeriodic)
  2273  	}
  2274  
  2275  	if !reflect.DeepEqual(nonPeriodic, outNonPeriodic) {
  2276  		t.Fatalf("bad: %#v %#v", nonPeriodic, outNonPeriodic)
  2277  	}
  2278  	if watchFired(ws) {
  2279  		t.Fatalf("bad")
  2280  	}
  2281  }
  2282  
  2283  func TestStateStore_JobsByScheduler(t *testing.T) {
  2284  	t.Parallel()
  2285  
  2286  	state := testStateStore(t)
  2287  	var serviceJobs []*structs.Job
  2288  	var sysJobs []*structs.Job
  2289  
  2290  	for i := 0; i < 10; i++ {
  2291  		job := mock.Job()
  2292  		serviceJobs = append(serviceJobs, job)
  2293  
  2294  		err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job)
  2295  		if err != nil {
  2296  			t.Fatalf("err: %v", err)
  2297  		}
  2298  	}
  2299  
  2300  	for i := 0; i < 10; i++ {
  2301  		job := mock.SystemJob()
  2302  		job.Status = structs.JobStatusRunning
  2303  		sysJobs = append(sysJobs, job)
  2304  
  2305  		err := state.UpsertJob(structs.MsgTypeTestSetup, 2000+uint64(i), job)
  2306  		if err != nil {
  2307  			t.Fatalf("err: %v", err)
  2308  		}
  2309  	}
  2310  
  2311  	ws := memdb.NewWatchSet()
  2312  	iter, err := state.JobsByScheduler(ws, "service")
  2313  	if err != nil {
  2314  		t.Fatalf("err: %v", err)
  2315  	}
  2316  
  2317  	var outService []*structs.Job
  2318  	for {
  2319  		raw := iter.Next()
  2320  		if raw == nil {
  2321  			break
  2322  		}
  2323  		outService = append(outService, raw.(*structs.Job))
  2324  	}
  2325  
  2326  	iter, err = state.JobsByScheduler(ws, "system")
  2327  	if err != nil {
  2328  		t.Fatalf("err: %v", err)
  2329  	}
  2330  
  2331  	var outSystem []*structs.Job
  2332  	for {
  2333  		raw := iter.Next()
  2334  		if raw == nil {
  2335  			break
  2336  		}
  2337  		outSystem = append(outSystem, raw.(*structs.Job))
  2338  	}
  2339  
  2340  	sort.Sort(JobIDSort(serviceJobs))
  2341  	sort.Sort(JobIDSort(sysJobs))
  2342  	sort.Sort(JobIDSort(outService))
  2343  	sort.Sort(JobIDSort(outSystem))
  2344  
  2345  	if !reflect.DeepEqual(serviceJobs, outService) {
  2346  		t.Fatalf("bad: %#v %#v", serviceJobs, outService)
  2347  	}
  2348  
  2349  	if !reflect.DeepEqual(sysJobs, outSystem) {
  2350  		t.Fatalf("bad: %#v %#v", sysJobs, outSystem)
  2351  	}
  2352  	if watchFired(ws) {
  2353  		t.Fatalf("bad")
  2354  	}
  2355  }
  2356  
  2357  func TestStateStore_JobsByGC(t *testing.T) {
  2358  	t.Parallel()
  2359  
  2360  	state := testStateStore(t)
  2361  	gc, nonGc := make(map[string]struct{}), make(map[string]struct{})
  2362  
  2363  	for i := 0; i < 20; i++ {
  2364  		var job *structs.Job
  2365  		if i%2 == 0 {
  2366  			job = mock.Job()
  2367  		} else {
  2368  			job = mock.PeriodicJob()
  2369  		}
  2370  		nonGc[job.ID] = struct{}{}
  2371  
  2372  		if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000+uint64(i), job); err != nil {
  2373  			t.Fatalf("err: %v", err)
  2374  		}
  2375  	}
  2376  
  2377  	for i := 0; i < 20; i += 2 {
  2378  		job := mock.Job()
  2379  		job.Type = structs.JobTypeBatch
  2380  		gc[job.ID] = struct{}{}
  2381  
  2382  		if err := state.UpsertJob(structs.MsgTypeTestSetup, 2000+uint64(i), job); err != nil {
  2383  			t.Fatalf("err: %v", err)
  2384  		}
  2385  
  2386  		// Create an eval for it
  2387  		eval := mock.Eval()
  2388  		eval.JobID = job.ID
  2389  		eval.Status = structs.EvalStatusComplete
  2390  		if err := state.UpsertEvals(structs.MsgTypeTestSetup, 2000+uint64(i+1), []*structs.Evaluation{eval}); err != nil {
  2391  			t.Fatalf("err: %v", err)
  2392  		}
  2393  
  2394  	}
  2395  
  2396  	ws := memdb.NewWatchSet()
  2397  	iter, err := state.JobsByGC(ws, true)
  2398  	if err != nil {
  2399  		t.Fatalf("err: %v", err)
  2400  	}
  2401  
  2402  	outGc := make(map[string]struct{})
  2403  	for i := iter.Next(); i != nil; i = iter.Next() {
  2404  		j := i.(*structs.Job)
  2405  		outGc[j.ID] = struct{}{}
  2406  	}
  2407  
  2408  	iter, err = state.JobsByGC(ws, false)
  2409  	if err != nil {
  2410  		t.Fatalf("err: %v", err)
  2411  	}
  2412  
  2413  	outNonGc := make(map[string]struct{})
  2414  	for i := iter.Next(); i != nil; i = iter.Next() {
  2415  		j := i.(*structs.Job)
  2416  		outNonGc[j.ID] = struct{}{}
  2417  	}
  2418  
  2419  	if !reflect.DeepEqual(gc, outGc) {
  2420  		t.Fatalf("bad: %#v %#v", gc, outGc)
  2421  	}
  2422  
  2423  	if !reflect.DeepEqual(nonGc, outNonGc) {
  2424  		t.Fatalf("bad: %#v %#v", nonGc, outNonGc)
  2425  	}
  2426  	if watchFired(ws) {
  2427  		t.Fatalf("bad")
  2428  	}
  2429  }
  2430  
  2431  func TestStateStore_RestoreJob(t *testing.T) {
  2432  	t.Parallel()
  2433  
  2434  	state := testStateStore(t)
  2435  	job := mock.Job()
  2436  
  2437  	restore, err := state.Restore()
  2438  	if err != nil {
  2439  		t.Fatalf("err: %v", err)
  2440  	}
  2441  
  2442  	err = restore.JobRestore(job)
  2443  	if err != nil {
  2444  		t.Fatalf("err: %v", err)
  2445  	}
  2446  	require.NoError(t, restore.Commit())
  2447  
  2448  	ws := memdb.NewWatchSet()
  2449  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  2450  	if err != nil {
  2451  		t.Fatalf("err: %v", err)
  2452  	}
  2453  
  2454  	if !reflect.DeepEqual(out, job) {
  2455  		t.Fatalf("Bad: %#v %#v", out, job)
  2456  	}
  2457  }
  2458  
  2459  func TestStateStore_UpsertPeriodicLaunch(t *testing.T) {
  2460  	t.Parallel()
  2461  
  2462  	state := testStateStore(t)
  2463  	job := mock.Job()
  2464  	launch := &structs.PeriodicLaunch{
  2465  		ID:        job.ID,
  2466  		Namespace: job.Namespace,
  2467  		Launch:    time.Now(),
  2468  	}
  2469  
  2470  	// Create a watchset so we can test that upsert fires the watch
  2471  	ws := memdb.NewWatchSet()
  2472  	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
  2473  		t.Fatalf("bad: %v", err)
  2474  	}
  2475  
  2476  	err := state.UpsertPeriodicLaunch(1000, launch)
  2477  	if err != nil {
  2478  		t.Fatalf("err: %v", err)
  2479  	}
  2480  
  2481  	if !watchFired(ws) {
  2482  		t.Fatalf("bad")
  2483  	}
  2484  
  2485  	ws = memdb.NewWatchSet()
  2486  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2487  	if err != nil {
  2488  		t.Fatalf("err: %v", err)
  2489  	}
  2490  	if out.CreateIndex != 1000 {
  2491  		t.Fatalf("bad: %#v", out)
  2492  	}
  2493  	if out.ModifyIndex != 1000 {
  2494  		t.Fatalf("bad: %#v", out)
  2495  	}
  2496  
  2497  	if !reflect.DeepEqual(launch, out) {
  2498  		t.Fatalf("bad: %#v %#v", job, out)
  2499  	}
  2500  
  2501  	index, err := state.Index("periodic_launch")
  2502  	if err != nil {
  2503  		t.Fatalf("err: %v", err)
  2504  	}
  2505  	if index != 1000 {
  2506  		t.Fatalf("bad: %d", index)
  2507  	}
  2508  
  2509  	if watchFired(ws) {
  2510  		t.Fatalf("bad")
  2511  	}
  2512  }
  2513  
  2514  func TestStateStore_UpdateUpsertPeriodicLaunch(t *testing.T) {
  2515  	t.Parallel()
  2516  
  2517  	state := testStateStore(t)
  2518  	job := mock.Job()
  2519  	launch := &structs.PeriodicLaunch{
  2520  		ID:        job.ID,
  2521  		Namespace: job.Namespace,
  2522  		Launch:    time.Now(),
  2523  	}
  2524  
  2525  	err := state.UpsertPeriodicLaunch(1000, launch)
  2526  	if err != nil {
  2527  		t.Fatalf("err: %v", err)
  2528  	}
  2529  
  2530  	// Create a watchset so we can test that upsert fires the watch
  2531  	ws := memdb.NewWatchSet()
  2532  	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
  2533  		t.Fatalf("bad: %v", err)
  2534  	}
  2535  
  2536  	launch2 := &structs.PeriodicLaunch{
  2537  		ID:        job.ID,
  2538  		Namespace: job.Namespace,
  2539  		Launch:    launch.Launch.Add(1 * time.Second),
  2540  	}
  2541  	err = state.UpsertPeriodicLaunch(1001, launch2)
  2542  	if err != nil {
  2543  		t.Fatalf("err: %v", err)
  2544  	}
  2545  
  2546  	if !watchFired(ws) {
  2547  		t.Fatalf("bad")
  2548  	}
  2549  
  2550  	ws = memdb.NewWatchSet()
  2551  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2552  	if err != nil {
  2553  		t.Fatalf("err: %v", err)
  2554  	}
  2555  	if out.CreateIndex != 1000 {
  2556  		t.Fatalf("bad: %#v", out)
  2557  	}
  2558  	if out.ModifyIndex != 1001 {
  2559  		t.Fatalf("bad: %#v", out)
  2560  	}
  2561  
  2562  	if !reflect.DeepEqual(launch2, out) {
  2563  		t.Fatalf("bad: %#v %#v", launch2, out)
  2564  	}
  2565  
  2566  	index, err := state.Index("periodic_launch")
  2567  	if err != nil {
  2568  		t.Fatalf("err: %v", err)
  2569  	}
  2570  	if index != 1001 {
  2571  		t.Fatalf("bad: %d", index)
  2572  	}
  2573  
  2574  	if watchFired(ws) {
  2575  		t.Fatalf("bad")
  2576  	}
  2577  }
  2578  
  2579  func TestStateStore_DeletePeriodicLaunch(t *testing.T) {
  2580  	t.Parallel()
  2581  
  2582  	state := testStateStore(t)
  2583  	job := mock.Job()
  2584  	launch := &structs.PeriodicLaunch{
  2585  		ID:        job.ID,
  2586  		Namespace: job.Namespace,
  2587  		Launch:    time.Now(),
  2588  	}
  2589  
  2590  	err := state.UpsertPeriodicLaunch(1000, launch)
  2591  	if err != nil {
  2592  		t.Fatalf("err: %v", err)
  2593  	}
  2594  
  2595  	// Create a watchset so we can test that delete fires the watch
  2596  	ws := memdb.NewWatchSet()
  2597  	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
  2598  		t.Fatalf("bad: %v", err)
  2599  	}
  2600  
  2601  	err = state.DeletePeriodicLaunch(1001, launch.Namespace, launch.ID)
  2602  	if err != nil {
  2603  		t.Fatalf("err: %v", err)
  2604  	}
  2605  
  2606  	if !watchFired(ws) {
  2607  		t.Fatalf("bad")
  2608  	}
  2609  
  2610  	ws = memdb.NewWatchSet()
  2611  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2612  	if err != nil {
  2613  		t.Fatalf("err: %v", err)
  2614  	}
  2615  
  2616  	if out != nil {
  2617  		t.Fatalf("bad: %#v %#v", job, out)
  2618  	}
  2619  
  2620  	index, err := state.Index("periodic_launch")
  2621  	if err != nil {
  2622  		t.Fatalf("err: %v", err)
  2623  	}
  2624  	if index != 1001 {
  2625  		t.Fatalf("bad: %d", index)
  2626  	}
  2627  
  2628  	if watchFired(ws) {
  2629  		t.Fatalf("bad")
  2630  	}
  2631  }
  2632  
  2633  func TestStateStore_PeriodicLaunches(t *testing.T) {
  2634  	t.Parallel()
  2635  
  2636  	state := testStateStore(t)
  2637  	var launches []*structs.PeriodicLaunch
  2638  
  2639  	for i := 0; i < 10; i++ {
  2640  		job := mock.Job()
  2641  		launch := &structs.PeriodicLaunch{
  2642  			ID:        job.ID,
  2643  			Namespace: job.Namespace,
  2644  			Launch:    time.Now(),
  2645  		}
  2646  		launches = append(launches, launch)
  2647  
  2648  		err := state.UpsertPeriodicLaunch(1000+uint64(i), launch)
  2649  		if err != nil {
  2650  			t.Fatalf("err: %v", err)
  2651  		}
  2652  	}
  2653  
  2654  	ws := memdb.NewWatchSet()
  2655  	iter, err := state.PeriodicLaunches(ws)
  2656  	if err != nil {
  2657  		t.Fatalf("err: %v", err)
  2658  	}
  2659  
  2660  	out := make(map[string]*structs.PeriodicLaunch, 10)
  2661  	for {
  2662  		raw := iter.Next()
  2663  		if raw == nil {
  2664  			break
  2665  		}
  2666  		launch := raw.(*structs.PeriodicLaunch)
  2667  		if _, ok := out[launch.ID]; ok {
  2668  			t.Fatalf("duplicate: %v", launch.ID)
  2669  		}
  2670  
  2671  		out[launch.ID] = launch
  2672  	}
  2673  
  2674  	for _, launch := range launches {
  2675  		l, ok := out[launch.ID]
  2676  		if !ok {
  2677  			t.Fatalf("bad %v", launch.ID)
  2678  		}
  2679  
  2680  		if !reflect.DeepEqual(launch, l) {
  2681  			t.Fatalf("bad: %#v %#v", launch, l)
  2682  		}
  2683  
  2684  		delete(out, launch.ID)
  2685  	}
  2686  
  2687  	if len(out) != 0 {
  2688  		t.Fatalf("leftover: %#v", out)
  2689  	}
  2690  
  2691  	if watchFired(ws) {
  2692  		t.Fatalf("bad")
  2693  	}
  2694  }
  2695  
  2696  func TestStateStore_RestorePeriodicLaunch(t *testing.T) {
  2697  	t.Parallel()
  2698  
  2699  	state := testStateStore(t)
  2700  	job := mock.Job()
  2701  	launch := &structs.PeriodicLaunch{
  2702  		ID:        job.ID,
  2703  		Namespace: job.Namespace,
  2704  		Launch:    time.Now(),
  2705  	}
  2706  
  2707  	restore, err := state.Restore()
  2708  	if err != nil {
  2709  		t.Fatalf("err: %v", err)
  2710  	}
  2711  
  2712  	err = restore.PeriodicLaunchRestore(launch)
  2713  	if err != nil {
  2714  		t.Fatalf("err: %v", err)
  2715  	}
  2716  	require.NoError(t, restore.Commit())
  2717  
  2718  	ws := memdb.NewWatchSet()
  2719  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2720  	if err != nil {
  2721  		t.Fatalf("err: %v", err)
  2722  	}
  2723  
  2724  	if !reflect.DeepEqual(out, launch) {
  2725  		t.Fatalf("Bad: %#v %#v", out, job)
  2726  	}
  2727  
  2728  	if watchFired(ws) {
  2729  		t.Fatalf("bad")
  2730  	}
  2731  }
  2732  
  2733  func TestStateStore_RestoreJobVersion(t *testing.T) {
  2734  	t.Parallel()
  2735  
  2736  	state := testStateStore(t)
  2737  	job := mock.Job()
  2738  
  2739  	restore, err := state.Restore()
  2740  	if err != nil {
  2741  		t.Fatalf("err: %v", err)
  2742  	}
  2743  
  2744  	err = restore.JobVersionRestore(job)
  2745  	if err != nil {
  2746  		t.Fatalf("err: %v", err)
  2747  	}
  2748  	require.NoError(t, restore.Commit())
  2749  
  2750  	ws := memdb.NewWatchSet()
  2751  	out, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, job.Version)
  2752  	if err != nil {
  2753  		t.Fatalf("err: %v", err)
  2754  	}
  2755  
  2756  	if !reflect.DeepEqual(out, job) {
  2757  		t.Fatalf("Bad: %#v %#v", out, job)
  2758  	}
  2759  
  2760  	if watchFired(ws) {
  2761  		t.Fatalf("bad")
  2762  	}
  2763  }
  2764  
  2765  func TestStateStore_RestoreDeployment(t *testing.T) {
  2766  	t.Parallel()
  2767  
  2768  	state := testStateStore(t)
  2769  	d := mock.Deployment()
  2770  
  2771  	restore, err := state.Restore()
  2772  	if err != nil {
  2773  		t.Fatalf("err: %v", err)
  2774  	}
  2775  
  2776  	err = restore.DeploymentRestore(d)
  2777  	if err != nil {
  2778  		t.Fatalf("err: %v", err)
  2779  	}
  2780  	require.NoError(t, restore.Commit())
  2781  
  2782  	ws := memdb.NewWatchSet()
  2783  	out, err := state.DeploymentByID(ws, d.ID)
  2784  	if err != nil {
  2785  		t.Fatalf("err: %v", err)
  2786  	}
  2787  
  2788  	if !reflect.DeepEqual(out, d) {
  2789  		t.Fatalf("Bad: %#v %#v", out, d)
  2790  	}
  2791  
  2792  	if watchFired(ws) {
  2793  		t.Fatalf("bad")
  2794  	}
  2795  }
  2796  
  2797  func TestStateStore_RestoreJobSummary(t *testing.T) {
  2798  	t.Parallel()
  2799  
  2800  	state := testStateStore(t)
  2801  	job := mock.Job()
  2802  	jobSummary := &structs.JobSummary{
  2803  		JobID:     job.ID,
  2804  		Namespace: job.Namespace,
  2805  		Summary: map[string]structs.TaskGroupSummary{
  2806  			"web": {
  2807  				Starting: 10,
  2808  			},
  2809  		},
  2810  	}
  2811  	restore, err := state.Restore()
  2812  	if err != nil {
  2813  		t.Fatalf("err: %v", err)
  2814  	}
  2815  
  2816  	err = restore.JobSummaryRestore(jobSummary)
  2817  	if err != nil {
  2818  		t.Fatalf("err: %v", err)
  2819  	}
  2820  	require.NoError(t, restore.Commit())
  2821  
  2822  	ws := memdb.NewWatchSet()
  2823  	out, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  2824  	if err != nil {
  2825  		t.Fatalf("err: %v", err)
  2826  	}
  2827  
  2828  	if !reflect.DeepEqual(out, jobSummary) {
  2829  		t.Fatalf("Bad: %#v %#v", out, jobSummary)
  2830  	}
  2831  }
  2832  
  2833  // TestStateStore_CSIVolume checks register, list and deregister for csi_volumes
  2834  func TestStateStore_CSIVolume(t *testing.T) {
  2835  	state := testStateStore(t)
  2836  	index := uint64(1000)
  2837  
  2838  	// Volume IDs
  2839  	vol0, vol1 := uuid.Generate(), uuid.Generate()
  2840  
  2841  	// Create a node running a healthy instance of the plugin
  2842  	node := mock.Node()
  2843  	pluginID := "minnie"
  2844  	alloc := mock.Alloc()
  2845  	alloc.DesiredStatus = "run"
  2846  	alloc.ClientStatus = "running"
  2847  	alloc.NodeID = node.ID
  2848  	alloc.Job.TaskGroups[0].Volumes = map[string]*structs.VolumeRequest{
  2849  		"foo": {
  2850  			Name:   "foo",
  2851  			Source: vol0,
  2852  			Type:   "csi",
  2853  		},
  2854  	}
  2855  
  2856  	node.CSINodePlugins = map[string]*structs.CSIInfo{
  2857  		pluginID: {
  2858  			PluginID:                 pluginID,
  2859  			AllocID:                  alloc.ID,
  2860  			Healthy:                  true,
  2861  			HealthDescription:        "healthy",
  2862  			RequiresControllerPlugin: false,
  2863  			RequiresTopologies:       false,
  2864  			NodeInfo: &structs.CSINodeInfo{
  2865  				ID:                      node.ID,
  2866  				MaxVolumes:              64,
  2867  				RequiresNodeStageVolume: true,
  2868  			},
  2869  		},
  2870  	}
  2871  
  2872  	index++
  2873  	err := state.UpsertNode(structs.MsgTypeTestSetup, index, node)
  2874  	require.NoError(t, err)
  2875  	defer state.DeleteNode(structs.MsgTypeTestSetup, 9999, []string{pluginID})
  2876  
  2877  	index++
  2878  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc})
  2879  	require.NoError(t, err)
  2880  
  2881  	ns := structs.DefaultNamespace
  2882  
  2883  	v0 := structs.NewCSIVolume("foo", index)
  2884  	v0.ID = vol0
  2885  	v0.Namespace = ns
  2886  	v0.PluginID = "minnie"
  2887  	v0.Schedulable = true
  2888  	v0.AccessMode = structs.CSIVolumeAccessModeMultiNodeSingleWriter
  2889  	v0.AttachmentMode = structs.CSIVolumeAttachmentModeFilesystem
  2890  
  2891  	index++
  2892  	v1 := structs.NewCSIVolume("foo", index)
  2893  	v1.ID = vol1
  2894  	v1.Namespace = ns
  2895  	v1.PluginID = "adam"
  2896  	v1.Schedulable = true
  2897  	v1.AccessMode = structs.CSIVolumeAccessModeMultiNodeSingleWriter
  2898  	v1.AttachmentMode = structs.CSIVolumeAttachmentModeFilesystem
  2899  
  2900  	index++
  2901  	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v0, v1})
  2902  	require.NoError(t, err)
  2903  
  2904  	// volume registration is idempotent, unless identies are changed
  2905  	index++
  2906  	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v0, v1})
  2907  	require.NoError(t, err)
  2908  
  2909  	index++
  2910  	v2 := v0.Copy()
  2911  	v2.PluginID = "new-id"
  2912  	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v2})
  2913  	require.Error(t, err, fmt.Sprintf("volume exists: %s", v0.ID))
  2914  
  2915  	ws := memdb.NewWatchSet()
  2916  	iter, err := state.CSIVolumesByNamespace(ws, ns)
  2917  	require.NoError(t, err)
  2918  
  2919  	slurp := func(iter memdb.ResultIterator) (vs []*structs.CSIVolume) {
  2920  		for {
  2921  			raw := iter.Next()
  2922  			if raw == nil {
  2923  				break
  2924  			}
  2925  			vol := raw.(*structs.CSIVolume)
  2926  			vs = append(vs, vol)
  2927  		}
  2928  		return vs
  2929  	}
  2930  
  2931  	vs := slurp(iter)
  2932  	require.Equal(t, 2, len(vs))
  2933  
  2934  	ws = memdb.NewWatchSet()
  2935  	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
  2936  	require.NoError(t, err)
  2937  	vs = slurp(iter)
  2938  	require.Equal(t, 1, len(vs))
  2939  
  2940  	ws = memdb.NewWatchSet()
  2941  	iter, err = state.CSIVolumesByNodeID(ws, node.ID)
  2942  	require.NoError(t, err)
  2943  	vs = slurp(iter)
  2944  	require.Equal(t, 1, len(vs))
  2945  
  2946  	// Allocs
  2947  	a0 := mock.Alloc()
  2948  	a1 := mock.Alloc()
  2949  	index++
  2950  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{a0, a1})
  2951  	require.NoError(t, err)
  2952  
  2953  	// Claims
  2954  	r := structs.CSIVolumeClaimRead
  2955  	w := structs.CSIVolumeClaimWrite
  2956  	u := structs.CSIVolumeClaimGC
  2957  	claim0 := &structs.CSIVolumeClaim{
  2958  		AllocationID: a0.ID,
  2959  		NodeID:       node.ID,
  2960  		Mode:         r,
  2961  	}
  2962  	claim1 := &structs.CSIVolumeClaim{
  2963  		AllocationID: a1.ID,
  2964  		NodeID:       node.ID,
  2965  		Mode:         w,
  2966  	}
  2967  
  2968  	index++
  2969  	err = state.CSIVolumeClaim(index, ns, vol0, claim0)
  2970  	require.NoError(t, err)
  2971  	index++
  2972  	err = state.CSIVolumeClaim(index, ns, vol0, claim1)
  2973  	require.NoError(t, err)
  2974  
  2975  	ws = memdb.NewWatchSet()
  2976  	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
  2977  	require.NoError(t, err)
  2978  	vs = slurp(iter)
  2979  	require.False(t, vs[0].WriteFreeClaims())
  2980  
  2981  	claim0.Mode = u
  2982  	err = state.CSIVolumeClaim(2, ns, vol0, claim0)
  2983  	require.NoError(t, err)
  2984  	ws = memdb.NewWatchSet()
  2985  	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
  2986  	require.NoError(t, err)
  2987  	vs = slurp(iter)
  2988  	require.True(t, vs[0].ReadSchedulable())
  2989  
  2990  	// registration is an error when the volume is in use
  2991  	index++
  2992  	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{v0})
  2993  	require.Error(t, err, fmt.Sprintf("volume exists: %s", vol0))
  2994  	// as is deregistration
  2995  	index++
  2996  	err = state.CSIVolumeDeregister(index, ns, []string{vol0}, false)
  2997  	require.Error(t, err, fmt.Sprintf("volume in use: %s", vol0))
  2998  
  2999  	// even if forced, because we have a non-terminal claim
  3000  	index++
  3001  	err = state.CSIVolumeDeregister(index, ns, []string{vol0}, true)
  3002  	require.Error(t, err, fmt.Sprintf("volume in use: %s", vol0))
  3003  
  3004  	// release claims to unblock deregister
  3005  	index++
  3006  	claim0.State = structs.CSIVolumeClaimStateReadyToFree
  3007  	err = state.CSIVolumeClaim(index, ns, vol0, claim0)
  3008  	require.NoError(t, err)
  3009  	index++
  3010  	claim1.Mode = u
  3011  	claim1.State = structs.CSIVolumeClaimStateReadyToFree
  3012  	err = state.CSIVolumeClaim(index, ns, vol0, claim1)
  3013  	require.NoError(t, err)
  3014  
  3015  	index++
  3016  	err = state.CSIVolumeDeregister(index, ns, []string{vol0}, false)
  3017  	require.NoError(t, err)
  3018  
  3019  	// List, now omitting the deregistered volume
  3020  	ws = memdb.NewWatchSet()
  3021  	iter, err = state.CSIVolumesByPluginID(ws, ns, "minnie")
  3022  	require.NoError(t, err)
  3023  	vs = slurp(iter)
  3024  	require.Equal(t, 0, len(vs))
  3025  
  3026  	ws = memdb.NewWatchSet()
  3027  	iter, err = state.CSIVolumesByNamespace(ws, ns)
  3028  	require.NoError(t, err)
  3029  	vs = slurp(iter)
  3030  	require.Equal(t, 1, len(vs))
  3031  }
  3032  
  3033  // TestStateStore_CSIPluginNodes uses node fingerprinting to create a plugin and update health
  3034  func TestStateStore_CSIPluginNodes(t *testing.T) {
  3035  	index := uint64(999)
  3036  	state := testStateStore(t)
  3037  	ws := memdb.NewWatchSet()
  3038  	plugID := "foo"
  3039  
  3040  	// Create Nomad client Nodes
  3041  	ns := []*structs.Node{mock.Node(), mock.Node()}
  3042  	for _, n := range ns {
  3043  		index++
  3044  		err := state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3045  		require.NoError(t, err)
  3046  	}
  3047  
  3048  	// Fingerprint a running controller plugin
  3049  	n0, _ := state.NodeByID(ws, ns[0].ID)
  3050  	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
  3051  		plugID: {
  3052  			PluginID:                 plugID,
  3053  			Healthy:                  true,
  3054  			UpdateTime:               time.Now(),
  3055  			RequiresControllerPlugin: true,
  3056  			RequiresTopologies:       false,
  3057  			ControllerInfo: &structs.CSIControllerInfo{
  3058  				SupportsReadOnlyAttach: true,
  3059  				SupportsListVolumes:    true,
  3060  			},
  3061  		},
  3062  	}
  3063  	index++
  3064  	err := state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3065  	require.NoError(t, err)
  3066  
  3067  	// Fingerprint two running node plugins
  3068  	for _, n := range ns[:] {
  3069  		n, _ := state.NodeByID(ws, n.ID)
  3070  		n.CSINodePlugins = map[string]*structs.CSIInfo{
  3071  			plugID: {
  3072  				PluginID:                 plugID,
  3073  				Healthy:                  true,
  3074  				UpdateTime:               time.Now(),
  3075  				RequiresControllerPlugin: true,
  3076  				RequiresTopologies:       false,
  3077  				NodeInfo:                 &structs.CSINodeInfo{},
  3078  			},
  3079  		}
  3080  		index++
  3081  		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3082  		require.NoError(t, err)
  3083  	}
  3084  
  3085  	plug, err := state.CSIPluginByID(ws, plugID)
  3086  	require.NoError(t, err)
  3087  	require.True(t, plug.ControllerRequired)
  3088  	require.Equal(t, 1, plug.ControllersHealthy, "controllers healthy")
  3089  	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
  3090  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3091  	require.Equal(t, 2, len(plug.Nodes), "nodes expected")
  3092  
  3093  	// Volume using the plugin
  3094  	index++
  3095  	vol := &structs.CSIVolume{
  3096  		ID:        uuid.Generate(),
  3097  		Namespace: structs.DefaultNamespace,
  3098  		PluginID:  plugID,
  3099  	}
  3100  	err = state.CSIVolumeRegister(index, []*structs.CSIVolume{vol})
  3101  	require.NoError(t, err)
  3102  
  3103  	vol, err = state.CSIVolumeByID(ws, structs.DefaultNamespace, vol.ID)
  3104  	require.NoError(t, err)
  3105  	require.True(t, vol.Schedulable, "volume should be schedulable")
  3106  
  3107  	// Controller is unhealthy
  3108  	n0, _ = state.NodeByID(ws, ns[0].ID)
  3109  	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
  3110  		plugID: {
  3111  			PluginID:                 plugID,
  3112  			Healthy:                  false,
  3113  			UpdateTime:               time.Now(),
  3114  			RequiresControllerPlugin: true,
  3115  			RequiresTopologies:       false,
  3116  			ControllerInfo: &structs.CSIControllerInfo{
  3117  				SupportsReadOnlyAttach: true,
  3118  				SupportsListVolumes:    true,
  3119  			},
  3120  		},
  3121  	}
  3122  
  3123  	index++
  3124  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3125  	require.NoError(t, err)
  3126  
  3127  	plug, err = state.CSIPluginByID(ws, plugID)
  3128  	require.NoError(t, err)
  3129  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3130  	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
  3131  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3132  	require.Equal(t, 2, len(plug.Nodes), "nodes expected")
  3133  
  3134  	vol, err = state.CSIVolumeByID(ws, structs.DefaultNamespace, vol.ID)
  3135  	require.NoError(t, err)
  3136  	require.False(t, vol.Schedulable, "volume should not be schedulable")
  3137  
  3138  	// Node plugin is removed
  3139  	n1, _ := state.NodeByID(ws, ns[1].ID)
  3140  	n1.CSINodePlugins = map[string]*structs.CSIInfo{}
  3141  	index++
  3142  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n1)
  3143  	require.NoError(t, err)
  3144  
  3145  	plug, err = state.CSIPluginByID(ws, plugID)
  3146  	require.NoError(t, err)
  3147  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3148  	require.Equal(t, 1, plug.NodesHealthy, "nodes healthy")
  3149  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3150  	require.Equal(t, 1, len(plug.Nodes), "nodes expected")
  3151  
  3152  	// Last node plugin is removed
  3153  	n0, _ = state.NodeByID(ws, ns[0].ID)
  3154  	n0.CSINodePlugins = map[string]*structs.CSIInfo{}
  3155  	index++
  3156  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3157  	require.NoError(t, err)
  3158  
  3159  	// Nodes plugins should be gone but controllers left
  3160  	plug, err = state.CSIPluginByID(ws, plugID)
  3161  	require.NoError(t, err)
  3162  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3163  	require.Equal(t, 0, plug.NodesHealthy, "nodes healthy")
  3164  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3165  	require.Equal(t, 0, len(plug.Nodes), "nodes expected")
  3166  
  3167  	// A node plugin is restored
  3168  	n0, _ = state.NodeByID(ws, n0.ID)
  3169  	n0.CSINodePlugins = map[string]*structs.CSIInfo{
  3170  		plugID: {
  3171  			PluginID:                 plugID,
  3172  			Healthy:                  true,
  3173  			UpdateTime:               time.Now(),
  3174  			RequiresControllerPlugin: true,
  3175  			RequiresTopologies:       false,
  3176  			NodeInfo:                 &structs.CSINodeInfo{},
  3177  		},
  3178  	}
  3179  	index++
  3180  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3181  	require.NoError(t, err)
  3182  
  3183  	// Nodes plugin should be replaced and healthy
  3184  	plug, err = state.CSIPluginByID(ws, plugID)
  3185  	require.NoError(t, err)
  3186  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3187  	require.Equal(t, 1, plug.NodesHealthy, "nodes healthy")
  3188  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3189  	require.Equal(t, 1, len(plug.Nodes), "nodes expected")
  3190  
  3191  	// Remove node again
  3192  	n0, _ = state.NodeByID(ws, ns[0].ID)
  3193  	n0.CSINodePlugins = map[string]*structs.CSIInfo{}
  3194  	index++
  3195  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3196  	require.NoError(t, err)
  3197  
  3198  	// Nodes plugins should be gone but controllers left
  3199  	plug, err = state.CSIPluginByID(ws, plugID)
  3200  	require.NoError(t, err)
  3201  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3202  	require.Equal(t, 0, plug.NodesHealthy, "nodes healthy")
  3203  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3204  	require.Equal(t, 0, len(plug.Nodes), "nodes expected")
  3205  
  3206  	// controller is removed
  3207  	n0, _ = state.NodeByID(ws, ns[0].ID)
  3208  	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{}
  3209  	index++
  3210  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3211  	require.NoError(t, err)
  3212  
  3213  	// Plugin has been removed entirely
  3214  	plug, err = state.CSIPluginByID(ws, plugID)
  3215  	require.NoError(t, err)
  3216  	require.Nil(t, plug)
  3217  
  3218  	// Volume still exists and is safe to query, but unschedulable
  3219  	vol, err = state.CSIVolumeByID(ws, structs.DefaultNamespace, vol.ID)
  3220  	require.NoError(t, err)
  3221  	require.False(t, vol.Schedulable)
  3222  }
  3223  
  3224  // TestStateStore_CSIPluginAllocUpdates tests the ordering
  3225  // interactions for CSI plugins between Nomad client node updates and
  3226  // allocation updates.
  3227  func TestStateStore_CSIPluginAllocUpdates(t *testing.T) {
  3228  	t.Parallel()
  3229  	index := uint64(999)
  3230  	state := testStateStore(t)
  3231  	ws := memdb.NewWatchSet()
  3232  
  3233  	n := mock.Node()
  3234  	index++
  3235  	err := state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3236  	require.NoError(t, err)
  3237  
  3238  	// (1) unhealthy fingerprint, then terminal alloc, then healthy node update
  3239  	plugID0 := "foo0"
  3240  
  3241  	alloc0 := mock.Alloc()
  3242  	alloc0.NodeID = n.ID
  3243  	alloc0.DesiredStatus = "run"
  3244  	alloc0.ClientStatus = "running"
  3245  	alloc0.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{ID: plugID0}
  3246  	index++
  3247  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc0})
  3248  	require.NoError(t, err)
  3249  
  3250  	n, _ = state.NodeByID(ws, n.ID)
  3251  	n.CSINodePlugins = map[string]*structs.CSIInfo{
  3252  		plugID0: {
  3253  			PluginID:                 plugID0,
  3254  			AllocID:                  alloc0.ID,
  3255  			Healthy:                  false,
  3256  			UpdateTime:               time.Now(),
  3257  			RequiresControllerPlugin: true,
  3258  			NodeInfo:                 &structs.CSINodeInfo{},
  3259  		},
  3260  	}
  3261  	index++
  3262  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3263  	require.NoError(t, err)
  3264  
  3265  	plug, err := state.CSIPluginByID(ws, plugID0)
  3266  	require.NoError(t, err)
  3267  	require.Nil(t, plug, "no plugin should exist: not yet healthy")
  3268  
  3269  	alloc0.DesiredStatus = "stopped"
  3270  	alloc0.ClientStatus = "complete"
  3271  	index++
  3272  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc0})
  3273  	require.NoError(t, err)
  3274  
  3275  	plug, err = state.CSIPluginByID(ws, plugID0)
  3276  	require.NoError(t, err)
  3277  	require.Nil(t, plug, "no plugin should exist: allocs never healthy")
  3278  
  3279  	n, _ = state.NodeByID(ws, n.ID)
  3280  	n.CSINodePlugins[plugID0].Healthy = true
  3281  	n.CSINodePlugins[plugID0].UpdateTime = time.Now()
  3282  	index++
  3283  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3284  	require.NoError(t, err)
  3285  
  3286  	plug, err = state.CSIPluginByID(ws, plugID0)
  3287  	require.NoError(t, err)
  3288  	require.NotNil(t, plug, "plugin should exist")
  3289  
  3290  	// (2) healthy fingerprint, then terminal alloc update
  3291  	plugID1 := "foo1"
  3292  
  3293  	alloc1 := mock.Alloc()
  3294  	n, _ = state.NodeByID(ws, n.ID)
  3295  	n.CSINodePlugins = map[string]*structs.CSIInfo{
  3296  		plugID1: {
  3297  			PluginID:                 plugID1,
  3298  			AllocID:                  alloc1.ID,
  3299  			Healthy:                  true,
  3300  			UpdateTime:               time.Now(),
  3301  			RequiresControllerPlugin: true,
  3302  			NodeInfo:                 &structs.CSINodeInfo{},
  3303  		},
  3304  	}
  3305  	index++
  3306  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3307  	require.NoError(t, err)
  3308  
  3309  	plug, err = state.CSIPluginByID(ws, plugID1)
  3310  	require.NoError(t, err)
  3311  	require.NotNil(t, plug, "plugin should exist")
  3312  
  3313  	alloc1.NodeID = n.ID
  3314  	alloc1.DesiredStatus = "stop"
  3315  	alloc1.ClientStatus = "complete"
  3316  	alloc1.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{ID: plugID1}
  3317  	index++
  3318  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1})
  3319  	require.NoError(t, err)
  3320  
  3321  	plug, err = state.CSIPluginByID(ws, plugID1)
  3322  	require.NoError(t, err)
  3323  	require.Nil(t, plug, "no plugin should exist: alloc became terminal")
  3324  
  3325  	// (3) terminal alloc update, then unhealthy fingerprint
  3326  	plugID2 := "foo2"
  3327  
  3328  	alloc2 := mock.Alloc()
  3329  	alloc2.NodeID = n.ID
  3330  	alloc2.DesiredStatus = "stop"
  3331  	alloc2.ClientStatus = "complete"
  3332  	alloc2.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{ID: plugID2}
  3333  	index++
  3334  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc2})
  3335  	require.NoError(t, err)
  3336  
  3337  	plug, err = state.CSIPluginByID(ws, plugID2)
  3338  	require.NoError(t, err)
  3339  	require.Nil(t, plug, "no plugin should exist: alloc became terminal")
  3340  
  3341  	n, _ = state.NodeByID(ws, n.ID)
  3342  	n.CSINodePlugins = map[string]*structs.CSIInfo{
  3343  		plugID2: {
  3344  			PluginID:                 plugID2,
  3345  			AllocID:                  alloc2.ID,
  3346  			Healthy:                  false,
  3347  			UpdateTime:               time.Now(),
  3348  			RequiresControllerPlugin: true,
  3349  			NodeInfo:                 &structs.CSINodeInfo{},
  3350  		},
  3351  	}
  3352  	index++
  3353  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3354  	require.NoError(t, err)
  3355  
  3356  	plug, err = state.CSIPluginByID(ws, plugID2)
  3357  	require.NoError(t, err)
  3358  	require.Nil(t, plug, "plugin should not exist: never became healthy")
  3359  
  3360  }
  3361  
  3362  // TestStateStore_CSIPluginMultiNodeUpdates tests the ordering
  3363  // interactions for CSI plugins between Nomad client node updates and
  3364  // allocation updates when multiple nodes are involved
  3365  func TestStateStore_CSIPluginMultiNodeUpdates(t *testing.T) {
  3366  	t.Parallel()
  3367  	index := uint64(999)
  3368  	state := testStateStore(t)
  3369  	ws := memdb.NewWatchSet()
  3370  
  3371  	var err error
  3372  
  3373  	// Create Nomad client Nodes
  3374  	ns := []*structs.Node{mock.Node(), mock.Node()}
  3375  	for _, n := range ns {
  3376  		index++
  3377  		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3378  		require.NoError(t, err)
  3379  	}
  3380  
  3381  	plugID := "foo"
  3382  	plugCfg := &structs.TaskCSIPluginConfig{ID: plugID}
  3383  
  3384  	// Fingerprint two running node plugins and their allocs; we'll
  3385  	// leave these in place for the test to ensure we don't GC the
  3386  	// plugin
  3387  	for _, n := range ns[:] {
  3388  		nAlloc := mock.Alloc()
  3389  		n, _ := state.NodeByID(ws, n.ID)
  3390  		n.CSINodePlugins = map[string]*structs.CSIInfo{
  3391  			plugID: {
  3392  				PluginID:                 plugID,
  3393  				AllocID:                  nAlloc.ID,
  3394  				Healthy:                  true,
  3395  				UpdateTime:               time.Now(),
  3396  				RequiresControllerPlugin: true,
  3397  				RequiresTopologies:       false,
  3398  				NodeInfo:                 &structs.CSINodeInfo{},
  3399  			},
  3400  		}
  3401  		index++
  3402  		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3403  		require.NoError(t, err)
  3404  
  3405  		nAlloc.NodeID = n.ID
  3406  		nAlloc.DesiredStatus = "run"
  3407  		nAlloc.ClientStatus = "running"
  3408  		nAlloc.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg
  3409  
  3410  		index++
  3411  		err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{nAlloc})
  3412  		require.NoError(t, err)
  3413  	}
  3414  
  3415  	// Fingerprint a running controller plugin
  3416  	alloc0 := mock.Alloc()
  3417  	n0, _ := state.NodeByID(ws, ns[0].ID)
  3418  	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
  3419  		plugID: {
  3420  			PluginID:                 plugID,
  3421  			AllocID:                  alloc0.ID,
  3422  			Healthy:                  true,
  3423  			UpdateTime:               time.Now(),
  3424  			RequiresControllerPlugin: true,
  3425  			RequiresTopologies:       false,
  3426  			ControllerInfo: &structs.CSIControllerInfo{
  3427  				SupportsReadOnlyAttach: true,
  3428  				SupportsListVolumes:    true,
  3429  			},
  3430  		},
  3431  	}
  3432  	index++
  3433  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3434  	require.NoError(t, err)
  3435  
  3436  	plug, err := state.CSIPluginByID(ws, plugID)
  3437  	require.NoError(t, err)
  3438  	require.Equal(t, 1, plug.ControllersHealthy, "controllers healthy")
  3439  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3440  	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
  3441  	require.Equal(t, 2, len(plug.Nodes), "nodes expected")
  3442  
  3443  	n1, _ := state.NodeByID(ws, ns[1].ID)
  3444  
  3445  	alloc0.NodeID = n0.ID
  3446  	alloc0.DesiredStatus = "stop"
  3447  	alloc0.ClientStatus = "complete"
  3448  	alloc0.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg
  3449  
  3450  	index++
  3451  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc0})
  3452  	require.NoError(t, err)
  3453  
  3454  	plug, err = state.CSIPluginByID(ws, plugID)
  3455  	require.NoError(t, err)
  3456  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3457  	require.Equal(t, 0, len(plug.Controllers), "controllers expected")
  3458  	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
  3459  	require.Equal(t, 2, len(plug.Nodes), "nodes expected")
  3460  
  3461  	alloc1 := mock.Alloc()
  3462  	alloc1.NodeID = n1.ID
  3463  	alloc1.DesiredStatus = "run"
  3464  	alloc1.ClientStatus = "running"
  3465  	alloc1.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg
  3466  
  3467  	index++
  3468  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc1})
  3469  	require.NoError(t, err)
  3470  
  3471  	plug, err = state.CSIPluginByID(ws, plugID)
  3472  	require.NoError(t, err)
  3473  	require.Equal(t, 0, plug.ControllersHealthy, "controllers healthy")
  3474  	require.Equal(t, 0, len(plug.Controllers), "controllers expected")
  3475  	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
  3476  	require.Equal(t, 2, len(plug.Nodes), "nodes expected")
  3477  
  3478  	n0, _ = state.NodeByID(ws, ns[0].ID)
  3479  	n0.CSIControllerPlugins = map[string]*structs.CSIInfo{
  3480  		plugID: {
  3481  			PluginID:                 plugID,
  3482  			AllocID:                  alloc0.ID,
  3483  			Healthy:                  false,
  3484  			UpdateTime:               time.Now(),
  3485  			RequiresControllerPlugin: true,
  3486  			RequiresTopologies:       false,
  3487  			ControllerInfo: &structs.CSIControllerInfo{
  3488  				SupportsReadOnlyAttach: true,
  3489  				SupportsListVolumes:    true,
  3490  			},
  3491  		},
  3492  	}
  3493  	index++
  3494  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n0)
  3495  	require.NoError(t, err)
  3496  
  3497  	n1.CSIControllerPlugins = map[string]*structs.CSIInfo{
  3498  		plugID: {
  3499  			PluginID:                 plugID,
  3500  			AllocID:                  alloc1.ID,
  3501  			Healthy:                  true,
  3502  			UpdateTime:               time.Now(),
  3503  			RequiresControllerPlugin: true,
  3504  			RequiresTopologies:       false,
  3505  			ControllerInfo: &structs.CSIControllerInfo{
  3506  				SupportsReadOnlyAttach: true,
  3507  				SupportsListVolumes:    true,
  3508  			},
  3509  		},
  3510  	}
  3511  	index++
  3512  	err = state.UpsertNode(structs.MsgTypeTestSetup, index, n1)
  3513  	require.NoError(t, err)
  3514  
  3515  	plug, err = state.CSIPluginByID(ws, plugID)
  3516  	require.NoError(t, err)
  3517  	require.True(t, plug.ControllerRequired)
  3518  	require.Equal(t, 1, plug.ControllersHealthy, "controllers healthy")
  3519  	require.Equal(t, 1, len(plug.Controllers), "controllers expected")
  3520  	require.Equal(t, 2, plug.NodesHealthy, "nodes healthy")
  3521  	require.Equal(t, 2, len(plug.Nodes), "nodes expected")
  3522  
  3523  }
  3524  
  3525  // TestStateStore_CSIPlugin_ConcurrentStop tests that concurrent allocation
  3526  // updates don't cause the count to drift unexpectedly or cause allocation
  3527  // update errors.
  3528  func TestStateStore_CSIPlugin_ConcurrentStop(t *testing.T) {
  3529  	t.Parallel()
  3530  	index := uint64(999)
  3531  	state := testStateStore(t)
  3532  	ws := memdb.NewWatchSet()
  3533  
  3534  	var err error
  3535  
  3536  	// Create Nomad client Nodes
  3537  	ns := []*structs.Node{mock.Node(), mock.Node(), mock.Node()}
  3538  	for _, n := range ns {
  3539  		index++
  3540  		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3541  		require.NoError(t, err)
  3542  	}
  3543  
  3544  	plugID := "foo"
  3545  	plugCfg := &structs.TaskCSIPluginConfig{ID: plugID}
  3546  
  3547  	allocs := []*structs.Allocation{}
  3548  
  3549  	// Fingerprint 3 running node plugins and their allocs
  3550  	for _, n := range ns[:] {
  3551  		alloc := mock.Alloc()
  3552  		n, _ := state.NodeByID(ws, n.ID)
  3553  		n.CSINodePlugins = map[string]*structs.CSIInfo{
  3554  			plugID: {
  3555  				PluginID:                 plugID,
  3556  				AllocID:                  alloc.ID,
  3557  				Healthy:                  true,
  3558  				UpdateTime:               time.Now(),
  3559  				RequiresControllerPlugin: true,
  3560  				RequiresTopologies:       false,
  3561  				NodeInfo:                 &structs.CSINodeInfo{},
  3562  			},
  3563  		}
  3564  		index++
  3565  		err = state.UpsertNode(structs.MsgTypeTestSetup, index, n)
  3566  		require.NoError(t, err)
  3567  
  3568  		alloc.NodeID = n.ID
  3569  		alloc.DesiredStatus = "run"
  3570  		alloc.ClientStatus = "running"
  3571  		alloc.Job.TaskGroups[0].Tasks[0].CSIPluginConfig = plugCfg
  3572  
  3573  		index++
  3574  		err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, []*structs.Allocation{alloc})
  3575  		require.NoError(t, err)
  3576  
  3577  		allocs = append(allocs, alloc)
  3578  	}
  3579  
  3580  	plug, err := state.CSIPluginByID(ws, plugID)
  3581  	require.NoError(t, err)
  3582  	require.Equal(t, 3, plug.NodesHealthy, "nodes healthy")
  3583  	require.Equal(t, 3, len(plug.Nodes), "nodes expected")
  3584  
  3585  	// stop all the allocs
  3586  	for _, alloc := range allocs {
  3587  		alloc.DesiredStatus = "stop"
  3588  		alloc.ClientStatus = "complete"
  3589  	}
  3590  
  3591  	// this is somewhat artificial b/c we get alloc updates from multiple
  3592  	// nodes concurrently but not in a single RPC call. But this guarantees
  3593  	// we'll trigger any nested transaction setup bugs
  3594  	index++
  3595  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, index, allocs)
  3596  	require.NoError(t, err)
  3597  
  3598  	plug, err = state.CSIPluginByID(ws, plugID)
  3599  	require.NoError(t, err)
  3600  	require.Nil(t, plug)
  3601  }
  3602  
  3603  func TestStateStore_CSIPluginJobs(t *testing.T) {
  3604  	s := testStateStore(t)
  3605  	index := uint64(1001)
  3606  
  3607  	controllerJob := mock.Job()
  3608  	controllerJob.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{
  3609  		ID:   "foo",
  3610  		Type: structs.CSIPluginTypeController,
  3611  	}
  3612  
  3613  	nodeJob := mock.Job()
  3614  	nodeJob.TaskGroups[0].Tasks[0].CSIPluginConfig = &structs.TaskCSIPluginConfig{
  3615  		ID:   "foo",
  3616  		Type: structs.CSIPluginTypeNode,
  3617  	}
  3618  
  3619  	err := s.UpsertJob(structs.MsgTypeTestSetup, index, controllerJob)
  3620  	require.NoError(t, err)
  3621  	index++
  3622  
  3623  	err = s.UpsertJob(structs.MsgTypeTestSetup, index, nodeJob)
  3624  	require.NoError(t, err)
  3625  	index++
  3626  
  3627  	// Get the plugin, and make better fake allocations for it
  3628  	ws := memdb.NewWatchSet()
  3629  	plug, err := s.CSIPluginByID(ws, "foo")
  3630  	require.NoError(t, err)
  3631  	index++
  3632  
  3633  	as := []*structs.Allocation{}
  3634  	for id, info := range plug.Controllers {
  3635  		as = append(as, &structs.Allocation{
  3636  			ID:        info.AllocID,
  3637  			Namespace: controllerJob.Namespace,
  3638  			JobID:     controllerJob.ID,
  3639  			Job:       controllerJob,
  3640  			TaskGroup: "web",
  3641  			EvalID:    uuid.Generate(),
  3642  			NodeID:    id,
  3643  		})
  3644  	}
  3645  	for id, info := range plug.Nodes {
  3646  		as = append(as, &structs.Allocation{
  3647  			ID:        info.AllocID,
  3648  			JobID:     nodeJob.ID,
  3649  			Namespace: nodeJob.Namespace,
  3650  			Job:       nodeJob,
  3651  			TaskGroup: "web",
  3652  			EvalID:    uuid.Generate(),
  3653  			NodeID:    id,
  3654  		})
  3655  	}
  3656  
  3657  	err = s.UpsertAllocs(structs.MsgTypeTestSetup, index, as)
  3658  	require.NoError(t, err)
  3659  	index++
  3660  
  3661  	// We use the summary to add
  3662  	err = s.ReconcileJobSummaries(index)
  3663  	require.NoError(t, err)
  3664  	index++
  3665  
  3666  	// Delete a job
  3667  	err = s.DeleteJob(index, controllerJob.Namespace, controllerJob.ID)
  3668  	require.NoError(t, err)
  3669  	index++
  3670  
  3671  	// plugin still exists
  3672  	plug, err = s.CSIPluginByID(ws, "foo")
  3673  	require.NoError(t, err)
  3674  	require.NotNil(t, plug)
  3675  	require.Equal(t, 0, len(plug.Controllers))
  3676  
  3677  	// Delete a job
  3678  	err = s.DeleteJob(index, nodeJob.Namespace, nodeJob.ID)
  3679  	require.NoError(t, err)
  3680  	index++
  3681  
  3682  	// plugin was collected
  3683  	plug, err = s.CSIPluginByID(ws, "foo")
  3684  	require.NoError(t, err)
  3685  	require.True(t, plug.IsEmpty())
  3686  }
  3687  
  3688  func TestStateStore_RestoreCSIPlugin(t *testing.T) {
  3689  	t.Parallel()
  3690  	require := require.New(t)
  3691  
  3692  	state := testStateStore(t)
  3693  	plugin := mock.CSIPlugin()
  3694  
  3695  	restore, err := state.Restore()
  3696  	require.NoError(err)
  3697  
  3698  	err = restore.CSIPluginRestore(plugin)
  3699  	require.NoError(err)
  3700  	require.NoError(restore.Commit())
  3701  
  3702  	ws := memdb.NewWatchSet()
  3703  	out, err := state.CSIPluginByID(ws, plugin.ID)
  3704  	require.NoError(err)
  3705  	require.EqualValues(out, plugin)
  3706  }
  3707  
  3708  func TestStateStore_RestoreCSIVolume(t *testing.T) {
  3709  	t.Parallel()
  3710  	require := require.New(t)
  3711  
  3712  	state := testStateStore(t)
  3713  	plugin := mock.CSIPlugin()
  3714  	volume := mock.CSIVolume(plugin)
  3715  
  3716  	restore, err := state.Restore()
  3717  	require.NoError(err)
  3718  
  3719  	err = restore.CSIVolumeRestore(volume)
  3720  	require.NoError(err)
  3721  	restore.Commit()
  3722  
  3723  	ws := memdb.NewWatchSet()
  3724  	out, err := state.CSIVolumeByID(ws, "default", volume.ID)
  3725  	require.NoError(err)
  3726  	require.EqualValues(out, volume)
  3727  }
  3728  
  3729  func TestStateStore_Indexes(t *testing.T) {
  3730  	t.Parallel()
  3731  
  3732  	state := testStateStore(t)
  3733  	node := mock.Node()
  3734  
  3735  	err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, node)
  3736  	if err != nil {
  3737  		t.Fatalf("err: %v", err)
  3738  	}
  3739  
  3740  	iter, err := state.Indexes()
  3741  	if err != nil {
  3742  		t.Fatalf("err: %v", err)
  3743  	}
  3744  
  3745  	var out []*IndexEntry
  3746  	for {
  3747  		raw := iter.Next()
  3748  		if raw == nil {
  3749  			break
  3750  		}
  3751  		out = append(out, raw.(*IndexEntry))
  3752  	}
  3753  
  3754  	expect := &IndexEntry{"nodes", 1000}
  3755  	if l := len(out); l < 1 {
  3756  		t.Fatalf("unexpected number of index entries: %v", pretty.Sprint(out))
  3757  	}
  3758  
  3759  	for _, index := range out {
  3760  		if index.Key != expect.Key {
  3761  			continue
  3762  		}
  3763  		if index.Value != expect.Value {
  3764  			t.Fatalf("bad index; got %d; want %d", index.Value, expect.Value)
  3765  		}
  3766  
  3767  		// We matched
  3768  		return
  3769  	}
  3770  
  3771  	t.Fatal("did not find expected index entry")
  3772  }
  3773  
  3774  func TestStateStore_LatestIndex(t *testing.T) {
  3775  	t.Parallel()
  3776  
  3777  	state := testStateStore(t)
  3778  
  3779  	if err := state.UpsertNode(structs.MsgTypeTestSetup, 1000, mock.Node()); err != nil {
  3780  		t.Fatalf("err: %v", err)
  3781  	}
  3782  
  3783  	exp := uint64(2000)
  3784  	if err := state.UpsertJob(structs.MsgTypeTestSetup, exp, mock.Job()); err != nil {
  3785  		t.Fatalf("err: %v", err)
  3786  	}
  3787  
  3788  	latest, err := state.LatestIndex()
  3789  	if err != nil {
  3790  		t.Fatalf("err: %v", err)
  3791  	}
  3792  
  3793  	if latest != exp {
  3794  		t.Fatalf("LatestIndex() returned %d; want %d", latest, exp)
  3795  	}
  3796  }
  3797  
  3798  func TestStateStore_RestoreIndex(t *testing.T) {
  3799  	t.Parallel()
  3800  
  3801  	state := testStateStore(t)
  3802  
  3803  	restore, err := state.Restore()
  3804  	if err != nil {
  3805  		t.Fatalf("err: %v", err)
  3806  	}
  3807  
  3808  	index := &IndexEntry{"jobs", 1000}
  3809  	err = restore.IndexRestore(index)
  3810  	if err != nil {
  3811  		t.Fatalf("err: %v", err)
  3812  	}
  3813  
  3814  	require.NoError(t, restore.Commit())
  3815  
  3816  	out, err := state.Index("jobs")
  3817  	if err != nil {
  3818  		t.Fatalf("err: %v", err)
  3819  	}
  3820  
  3821  	if out != 1000 {
  3822  		t.Fatalf("Bad: %#v %#v", out, 1000)
  3823  	}
  3824  }
  3825  
  3826  func TestStateStore_UpsertEvals_Eval(t *testing.T) {
  3827  	t.Parallel()
  3828  
  3829  	state := testStateStore(t)
  3830  	eval := mock.Eval()
  3831  
  3832  	// Create a watchset so we can test that upsert fires the watch
  3833  	ws := memdb.NewWatchSet()
  3834  	if _, err := state.EvalByID(ws, eval.ID); err != nil {
  3835  		t.Fatalf("bad: %v", err)
  3836  	}
  3837  
  3838  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval})
  3839  	if err != nil {
  3840  		t.Fatalf("err: %v", err)
  3841  	}
  3842  
  3843  	if !watchFired(ws) {
  3844  		t.Fatalf("bad")
  3845  	}
  3846  
  3847  	ws = memdb.NewWatchSet()
  3848  	out, err := state.EvalByID(ws, eval.ID)
  3849  	if err != nil {
  3850  		t.Fatalf("err: %v", err)
  3851  	}
  3852  
  3853  	if !reflect.DeepEqual(eval, out) {
  3854  		t.Fatalf("bad: %#v %#v", eval, out)
  3855  	}
  3856  
  3857  	index, err := state.Index("evals")
  3858  	if err != nil {
  3859  		t.Fatalf("err: %v", err)
  3860  	}
  3861  	if index != 1000 {
  3862  		t.Fatalf("bad: %d", index)
  3863  	}
  3864  
  3865  	if watchFired(ws) {
  3866  		t.Fatalf("bad")
  3867  	}
  3868  }
  3869  
  3870  func TestStateStore_UpsertEvals_CancelBlocked(t *testing.T) {
  3871  	t.Parallel()
  3872  
  3873  	state := testStateStore(t)
  3874  
  3875  	// Create two blocked evals for the same job
  3876  	j := "test-job"
  3877  	b1, b2 := mock.Eval(), mock.Eval()
  3878  	b1.JobID = j
  3879  	b1.Status = structs.EvalStatusBlocked
  3880  	b2.JobID = j
  3881  	b2.Status = structs.EvalStatusBlocked
  3882  
  3883  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 999, []*structs.Evaluation{b1, b2})
  3884  	if err != nil {
  3885  		t.Fatalf("err: %v", err)
  3886  	}
  3887  
  3888  	// Create one complete and successful eval for the job
  3889  	eval := mock.Eval()
  3890  	eval.JobID = j
  3891  	eval.Status = structs.EvalStatusComplete
  3892  
  3893  	// Create a watchset so we can test that the upsert of the complete eval
  3894  	// fires the watch
  3895  	ws := memdb.NewWatchSet()
  3896  	if _, err := state.EvalByID(ws, b1.ID); err != nil {
  3897  		t.Fatalf("bad: %v", err)
  3898  	}
  3899  	if _, err := state.EvalByID(ws, b2.ID); err != nil {
  3900  		t.Fatalf("bad: %v", err)
  3901  	}
  3902  
  3903  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval}); err != nil {
  3904  		t.Fatalf("err: %v", err)
  3905  	}
  3906  
  3907  	if !watchFired(ws) {
  3908  		t.Fatalf("bad")
  3909  	}
  3910  
  3911  	ws = memdb.NewWatchSet()
  3912  	out, err := state.EvalByID(ws, eval.ID)
  3913  	if err != nil {
  3914  		t.Fatalf("err: %v", err)
  3915  	}
  3916  
  3917  	if !reflect.DeepEqual(eval, out) {
  3918  		t.Fatalf("bad: %#v %#v", eval, out)
  3919  	}
  3920  
  3921  	index, err := state.Index("evals")
  3922  	if err != nil {
  3923  		t.Fatalf("err: %v", err)
  3924  	}
  3925  	if index != 1000 {
  3926  		t.Fatalf("bad: %d", index)
  3927  	}
  3928  
  3929  	// Get b1/b2 and check they are cancelled
  3930  	out1, err := state.EvalByID(ws, b1.ID)
  3931  	if err != nil {
  3932  		t.Fatalf("err: %v", err)
  3933  	}
  3934  
  3935  	out2, err := state.EvalByID(ws, b2.ID)
  3936  	if err != nil {
  3937  		t.Fatalf("err: %v", err)
  3938  	}
  3939  
  3940  	if out1.Status != structs.EvalStatusCancelled || out2.Status != structs.EvalStatusCancelled {
  3941  		t.Fatalf("bad: %#v %#v", out1, out2)
  3942  	}
  3943  
  3944  	if watchFired(ws) {
  3945  		t.Fatalf("bad")
  3946  	}
  3947  }
  3948  
  3949  func TestStateStore_Update_UpsertEvals_Eval(t *testing.T) {
  3950  	t.Parallel()
  3951  
  3952  	state := testStateStore(t)
  3953  	eval := mock.Eval()
  3954  
  3955  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval})
  3956  	if err != nil {
  3957  		t.Fatalf("err: %v", err)
  3958  	}
  3959  
  3960  	// Create a watchset so we can test that delete fires the watch
  3961  	ws := memdb.NewWatchSet()
  3962  	ws2 := memdb.NewWatchSet()
  3963  	if _, err := state.EvalByID(ws, eval.ID); err != nil {
  3964  		t.Fatalf("bad: %v", err)
  3965  	}
  3966  
  3967  	if _, err := state.EvalsByJob(ws2, eval.Namespace, eval.JobID); err != nil {
  3968  		t.Fatalf("bad: %v", err)
  3969  	}
  3970  
  3971  	eval2 := mock.Eval()
  3972  	eval2.ID = eval.ID
  3973  	eval2.JobID = eval.JobID
  3974  	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1001, []*structs.Evaluation{eval2})
  3975  	if err != nil {
  3976  		t.Fatalf("err: %v", err)
  3977  	}
  3978  
  3979  	if !watchFired(ws) {
  3980  		t.Fatalf("bad")
  3981  	}
  3982  	if !watchFired(ws2) {
  3983  		t.Fatalf("bad")
  3984  	}
  3985  
  3986  	ws = memdb.NewWatchSet()
  3987  	out, err := state.EvalByID(ws, eval.ID)
  3988  	if err != nil {
  3989  		t.Fatalf("err: %v", err)
  3990  	}
  3991  
  3992  	if !reflect.DeepEqual(eval2, out) {
  3993  		t.Fatalf("bad: %#v %#v", eval2, out)
  3994  	}
  3995  
  3996  	if out.CreateIndex != 1000 {
  3997  		t.Fatalf("bad: %#v", out)
  3998  	}
  3999  	if out.ModifyIndex != 1001 {
  4000  		t.Fatalf("bad: %#v", out)
  4001  	}
  4002  
  4003  	index, err := state.Index("evals")
  4004  	if err != nil {
  4005  		t.Fatalf("err: %v", err)
  4006  	}
  4007  	if index != 1001 {
  4008  		t.Fatalf("bad: %d", index)
  4009  	}
  4010  
  4011  	if watchFired(ws) {
  4012  		t.Fatalf("bad")
  4013  	}
  4014  }
  4015  
  4016  func TestStateStore_UpsertEvals_Eval_ChildJob(t *testing.T) {
  4017  	t.Parallel()
  4018  
  4019  	state := testStateStore(t)
  4020  
  4021  	parent := mock.Job()
  4022  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
  4023  		t.Fatalf("err: %v", err)
  4024  	}
  4025  
  4026  	child := mock.Job()
  4027  	child.Status = ""
  4028  	child.ParentID = parent.ID
  4029  
  4030  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
  4031  		t.Fatalf("err: %v", err)
  4032  	}
  4033  
  4034  	eval := mock.Eval()
  4035  	eval.Status = structs.EvalStatusComplete
  4036  	eval.JobID = child.ID
  4037  
  4038  	// Create watchsets so we can test that upsert fires the watch
  4039  	ws := memdb.NewWatchSet()
  4040  	ws2 := memdb.NewWatchSet()
  4041  	ws3 := memdb.NewWatchSet()
  4042  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  4043  		t.Fatalf("bad: %v", err)
  4044  	}
  4045  	if _, err := state.EvalByID(ws2, eval.ID); err != nil {
  4046  		t.Fatalf("bad: %v", err)
  4047  	}
  4048  	if _, err := state.EvalsByJob(ws3, eval.Namespace, eval.JobID); err != nil {
  4049  		t.Fatalf("bad: %v", err)
  4050  	}
  4051  
  4052  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval})
  4053  	if err != nil {
  4054  		t.Fatalf("err: %v", err)
  4055  	}
  4056  
  4057  	if !watchFired(ws) {
  4058  		t.Fatalf("bad")
  4059  	}
  4060  	if !watchFired(ws2) {
  4061  		t.Fatalf("bad")
  4062  	}
  4063  	if !watchFired(ws3) {
  4064  		t.Fatalf("bad")
  4065  	}
  4066  
  4067  	ws = memdb.NewWatchSet()
  4068  	out, err := state.EvalByID(ws, eval.ID)
  4069  	if err != nil {
  4070  		t.Fatalf("err: %v", err)
  4071  	}
  4072  
  4073  	if !reflect.DeepEqual(eval, out) {
  4074  		t.Fatalf("bad: %#v %#v", eval, out)
  4075  	}
  4076  
  4077  	index, err := state.Index("evals")
  4078  	if err != nil {
  4079  		t.Fatalf("err: %v", err)
  4080  	}
  4081  	if index != 1000 {
  4082  		t.Fatalf("bad: %d", index)
  4083  	}
  4084  
  4085  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  4086  	if err != nil {
  4087  		t.Fatalf("err: %v", err)
  4088  	}
  4089  	if summary == nil {
  4090  		t.Fatalf("nil summary")
  4091  	}
  4092  	if summary.JobID != parent.ID {
  4093  		t.Fatalf("bad summary id: %v", parent.ID)
  4094  	}
  4095  	if summary.Children == nil {
  4096  		t.Fatalf("nil children summary")
  4097  	}
  4098  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  4099  		t.Fatalf("bad children summary: %v", summary.Children)
  4100  	}
  4101  
  4102  	if watchFired(ws) {
  4103  		t.Fatalf("bad")
  4104  	}
  4105  }
  4106  
  4107  func TestStateStore_DeleteEval_Eval(t *testing.T) {
  4108  	t.Parallel()
  4109  
  4110  	state := testStateStore(t)
  4111  	eval1 := mock.Eval()
  4112  	eval2 := mock.Eval()
  4113  	alloc1 := mock.Alloc()
  4114  	alloc2 := mock.Alloc()
  4115  
  4116  	// Create watchsets so we can test that upsert fires the watch
  4117  	watches := make([]memdb.WatchSet, 12)
  4118  	for i := 0; i < 12; i++ {
  4119  		watches[i] = memdb.NewWatchSet()
  4120  	}
  4121  	if _, err := state.EvalByID(watches[0], eval1.ID); err != nil {
  4122  		t.Fatalf("bad: %v", err)
  4123  	}
  4124  	if _, err := state.EvalByID(watches[1], eval2.ID); err != nil {
  4125  		t.Fatalf("bad: %v", err)
  4126  	}
  4127  	if _, err := state.EvalsByJob(watches[2], eval1.Namespace, eval1.JobID); err != nil {
  4128  		t.Fatalf("bad: %v", err)
  4129  	}
  4130  	if _, err := state.EvalsByJob(watches[3], eval2.Namespace, eval2.JobID); err != nil {
  4131  		t.Fatalf("bad: %v", err)
  4132  	}
  4133  	if _, err := state.AllocByID(watches[4], alloc1.ID); err != nil {
  4134  		t.Fatalf("bad: %v", err)
  4135  	}
  4136  	if _, err := state.AllocByID(watches[5], alloc2.ID); err != nil {
  4137  		t.Fatalf("bad: %v", err)
  4138  	}
  4139  	if _, err := state.AllocsByEval(watches[6], alloc1.EvalID); err != nil {
  4140  		t.Fatalf("bad: %v", err)
  4141  	}
  4142  	if _, err := state.AllocsByEval(watches[7], alloc2.EvalID); err != nil {
  4143  		t.Fatalf("bad: %v", err)
  4144  	}
  4145  	if _, err := state.AllocsByJob(watches[8], alloc1.Namespace, alloc1.JobID, false); err != nil {
  4146  		t.Fatalf("bad: %v", err)
  4147  	}
  4148  	if _, err := state.AllocsByJob(watches[9], alloc2.Namespace, alloc2.JobID, false); err != nil {
  4149  		t.Fatalf("bad: %v", err)
  4150  	}
  4151  	if _, err := state.AllocsByNode(watches[10], alloc1.NodeID); err != nil {
  4152  		t.Fatalf("bad: %v", err)
  4153  	}
  4154  	if _, err := state.AllocsByNode(watches[11], alloc2.NodeID); err != nil {
  4155  		t.Fatalf("bad: %v", err)
  4156  	}
  4157  
  4158  	state.UpsertJobSummary(900, mock.JobSummary(eval1.JobID))
  4159  	state.UpsertJobSummary(901, mock.JobSummary(eval2.JobID))
  4160  	state.UpsertJobSummary(902, mock.JobSummary(alloc1.JobID))
  4161  	state.UpsertJobSummary(903, mock.JobSummary(alloc2.JobID))
  4162  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval1, eval2})
  4163  	if err != nil {
  4164  		t.Fatalf("err: %v", err)
  4165  	}
  4166  
  4167  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc1, alloc2})
  4168  	if err != nil {
  4169  		t.Fatalf("err: %v", err)
  4170  	}
  4171  
  4172  	err = state.DeleteEval(1002, []string{eval1.ID, eval2.ID}, []string{alloc1.ID, alloc2.ID})
  4173  	if err != nil {
  4174  		t.Fatalf("err: %v", err)
  4175  	}
  4176  
  4177  	for i, ws := range watches {
  4178  		if !watchFired(ws) {
  4179  			t.Fatalf("bad %d", i)
  4180  		}
  4181  	}
  4182  
  4183  	ws := memdb.NewWatchSet()
  4184  	out, err := state.EvalByID(ws, eval1.ID)
  4185  	if err != nil {
  4186  		t.Fatalf("err: %v", err)
  4187  	}
  4188  
  4189  	if out != nil {
  4190  		t.Fatalf("bad: %#v %#v", eval1, out)
  4191  	}
  4192  
  4193  	out, err = state.EvalByID(ws, eval2.ID)
  4194  	if err != nil {
  4195  		t.Fatalf("err: %v", err)
  4196  	}
  4197  
  4198  	if out != nil {
  4199  		t.Fatalf("bad: %#v %#v", eval1, out)
  4200  	}
  4201  
  4202  	outA, err := state.AllocByID(ws, alloc1.ID)
  4203  	if err != nil {
  4204  		t.Fatalf("err: %v", err)
  4205  	}
  4206  
  4207  	if out != nil {
  4208  		t.Fatalf("bad: %#v %#v", alloc1, outA)
  4209  	}
  4210  
  4211  	outA, err = state.AllocByID(ws, alloc2.ID)
  4212  	if err != nil {
  4213  		t.Fatalf("err: %v", err)
  4214  	}
  4215  
  4216  	if out != nil {
  4217  		t.Fatalf("bad: %#v %#v", alloc1, outA)
  4218  	}
  4219  
  4220  	index, err := state.Index("evals")
  4221  	if err != nil {
  4222  		t.Fatalf("err: %v", err)
  4223  	}
  4224  	if index != 1002 {
  4225  		t.Fatalf("bad: %d", index)
  4226  	}
  4227  
  4228  	index, err = state.Index("allocs")
  4229  	if err != nil {
  4230  		t.Fatalf("err: %v", err)
  4231  	}
  4232  	if index != 1002 {
  4233  		t.Fatalf("bad: %d", index)
  4234  	}
  4235  
  4236  	if watchFired(ws) {
  4237  		t.Fatalf("bad")
  4238  	}
  4239  }
  4240  
  4241  func TestStateStore_DeleteEval_ChildJob(t *testing.T) {
  4242  	t.Parallel()
  4243  
  4244  	state := testStateStore(t)
  4245  
  4246  	parent := mock.Job()
  4247  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
  4248  		t.Fatalf("err: %v", err)
  4249  	}
  4250  
  4251  	child := mock.Job()
  4252  	child.Status = ""
  4253  	child.ParentID = parent.ID
  4254  
  4255  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
  4256  		t.Fatalf("err: %v", err)
  4257  	}
  4258  
  4259  	eval1 := mock.Eval()
  4260  	eval1.JobID = child.ID
  4261  	alloc1 := mock.Alloc()
  4262  	alloc1.JobID = child.ID
  4263  
  4264  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval1})
  4265  	if err != nil {
  4266  		t.Fatalf("err: %v", err)
  4267  	}
  4268  
  4269  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc1})
  4270  	if err != nil {
  4271  		t.Fatalf("err: %v", err)
  4272  	}
  4273  
  4274  	// Create watchsets so we can test that delete fires the watch
  4275  	ws := memdb.NewWatchSet()
  4276  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  4277  		t.Fatalf("bad: %v", err)
  4278  	}
  4279  
  4280  	err = state.DeleteEval(1002, []string{eval1.ID}, []string{alloc1.ID})
  4281  	if err != nil {
  4282  		t.Fatalf("err: %v", err)
  4283  	}
  4284  
  4285  	if !watchFired(ws) {
  4286  		t.Fatalf("bad")
  4287  	}
  4288  
  4289  	ws = memdb.NewWatchSet()
  4290  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  4291  	if err != nil {
  4292  		t.Fatalf("err: %v", err)
  4293  	}
  4294  	if summary == nil {
  4295  		t.Fatalf("nil summary")
  4296  	}
  4297  	if summary.JobID != parent.ID {
  4298  		t.Fatalf("bad summary id: %v", parent.ID)
  4299  	}
  4300  	if summary.Children == nil {
  4301  		t.Fatalf("nil children summary")
  4302  	}
  4303  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  4304  		t.Fatalf("bad children summary: %v", summary.Children)
  4305  	}
  4306  
  4307  	if watchFired(ws) {
  4308  		t.Fatalf("bad")
  4309  	}
  4310  }
  4311  
  4312  func TestStateStore_EvalsByJob(t *testing.T) {
  4313  	t.Parallel()
  4314  
  4315  	state := testStateStore(t)
  4316  
  4317  	eval1 := mock.Eval()
  4318  	eval2 := mock.Eval()
  4319  	eval2.JobID = eval1.JobID
  4320  	eval3 := mock.Eval()
  4321  	evals := []*structs.Evaluation{eval1, eval2}
  4322  
  4323  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, evals)
  4324  	if err != nil {
  4325  		t.Fatalf("err: %v", err)
  4326  	}
  4327  	err = state.UpsertEvals(structs.MsgTypeTestSetup, 1001, []*structs.Evaluation{eval3})
  4328  	if err != nil {
  4329  		t.Fatalf("err: %v", err)
  4330  	}
  4331  
  4332  	ws := memdb.NewWatchSet()
  4333  	out, err := state.EvalsByJob(ws, eval1.Namespace, eval1.JobID)
  4334  	if err != nil {
  4335  		t.Fatalf("err: %v", err)
  4336  	}
  4337  
  4338  	sort.Sort(EvalIDSort(evals))
  4339  	sort.Sort(EvalIDSort(out))
  4340  
  4341  	if !reflect.DeepEqual(evals, out) {
  4342  		t.Fatalf("bad: %#v %#v", evals, out)
  4343  	}
  4344  
  4345  	if watchFired(ws) {
  4346  		t.Fatalf("bad")
  4347  	}
  4348  }
  4349  
  4350  func TestStateStore_Evals(t *testing.T) {
  4351  	t.Parallel()
  4352  
  4353  	state := testStateStore(t)
  4354  	var evals []*structs.Evaluation
  4355  
  4356  	for i := 0; i < 10; i++ {
  4357  		eval := mock.Eval()
  4358  		evals = append(evals, eval)
  4359  
  4360  		err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000+uint64(i), []*structs.Evaluation{eval})
  4361  		if err != nil {
  4362  			t.Fatalf("err: %v", err)
  4363  		}
  4364  	}
  4365  
  4366  	ws := memdb.NewWatchSet()
  4367  	iter, err := state.Evals(ws)
  4368  	if err != nil {
  4369  		t.Fatalf("err: %v", err)
  4370  	}
  4371  
  4372  	var out []*structs.Evaluation
  4373  	for {
  4374  		raw := iter.Next()
  4375  		if raw == nil {
  4376  			break
  4377  		}
  4378  		out = append(out, raw.(*structs.Evaluation))
  4379  	}
  4380  
  4381  	sort.Sort(EvalIDSort(evals))
  4382  	sort.Sort(EvalIDSort(out))
  4383  
  4384  	if !reflect.DeepEqual(evals, out) {
  4385  		t.Fatalf("bad: %#v %#v", evals, out)
  4386  	}
  4387  
  4388  	if watchFired(ws) {
  4389  		t.Fatalf("bad")
  4390  	}
  4391  }
  4392  
  4393  func TestStateStore_EvalsByIDPrefix(t *testing.T) {
  4394  	t.Parallel()
  4395  
  4396  	state := testStateStore(t)
  4397  	var evals []*structs.Evaluation
  4398  
  4399  	ids := []string{
  4400  		"aaaaaaaa-7bfb-395d-eb95-0685af2176b2",
  4401  		"aaaaaaab-7bfb-395d-eb95-0685af2176b2",
  4402  		"aaaaaabb-7bfb-395d-eb95-0685af2176b2",
  4403  		"aaaaabbb-7bfb-395d-eb95-0685af2176b2",
  4404  		"aaaabbbb-7bfb-395d-eb95-0685af2176b2",
  4405  		"aaabbbbb-7bfb-395d-eb95-0685af2176b2",
  4406  		"aabbbbbb-7bfb-395d-eb95-0685af2176b2",
  4407  		"abbbbbbb-7bfb-395d-eb95-0685af2176b2",
  4408  		"bbbbbbbb-7bfb-395d-eb95-0685af2176b2",
  4409  	}
  4410  	for i := 0; i < 9; i++ {
  4411  		eval := mock.Eval()
  4412  		eval.ID = ids[i]
  4413  		evals = append(evals, eval)
  4414  	}
  4415  
  4416  	err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, evals)
  4417  	if err != nil {
  4418  		t.Fatalf("err: %v", err)
  4419  	}
  4420  
  4421  	ws := memdb.NewWatchSet()
  4422  	iter, err := state.EvalsByIDPrefix(ws, structs.DefaultNamespace, "aaaa")
  4423  	if err != nil {
  4424  		t.Fatalf("err: %v", err)
  4425  	}
  4426  
  4427  	gatherEvals := func(iter memdb.ResultIterator) []*structs.Evaluation {
  4428  		var evals []*structs.Evaluation
  4429  		for {
  4430  			raw := iter.Next()
  4431  			if raw == nil {
  4432  				break
  4433  			}
  4434  			evals = append(evals, raw.(*structs.Evaluation))
  4435  		}
  4436  		return evals
  4437  	}
  4438  
  4439  	out := gatherEvals(iter)
  4440  	if len(out) != 5 {
  4441  		t.Fatalf("bad: expected five evaluations, got: %#v", out)
  4442  	}
  4443  
  4444  	sort.Sort(EvalIDSort(evals))
  4445  
  4446  	for index, eval := range out {
  4447  		if ids[index] != eval.ID {
  4448  			t.Fatalf("bad: got unexpected id: %s", eval.ID)
  4449  		}
  4450  	}
  4451  
  4452  	iter, err = state.EvalsByIDPrefix(ws, structs.DefaultNamespace, "b-a7bfb")
  4453  	if err != nil {
  4454  		t.Fatalf("err: %v", err)
  4455  	}
  4456  
  4457  	out = gatherEvals(iter)
  4458  	if len(out) != 0 {
  4459  		t.Fatalf("bad: unexpected zero evaluations, got: %#v", out)
  4460  	}
  4461  
  4462  	if watchFired(ws) {
  4463  		t.Fatalf("bad")
  4464  	}
  4465  }
  4466  
  4467  func TestStateStore_RestoreEval(t *testing.T) {
  4468  	t.Parallel()
  4469  
  4470  	state := testStateStore(t)
  4471  	eval := mock.Eval()
  4472  
  4473  	restore, err := state.Restore()
  4474  	if err != nil {
  4475  		t.Fatalf("err: %v", err)
  4476  	}
  4477  
  4478  	err = restore.EvalRestore(eval)
  4479  	if err != nil {
  4480  		t.Fatalf("err: %v", err)
  4481  	}
  4482  	require.NoError(t, restore.Commit())
  4483  
  4484  	ws := memdb.NewWatchSet()
  4485  	out, err := state.EvalByID(ws, eval.ID)
  4486  	if err != nil {
  4487  		t.Fatalf("err: %v", err)
  4488  	}
  4489  
  4490  	if !reflect.DeepEqual(out, eval) {
  4491  		t.Fatalf("Bad: %#v %#v", out, eval)
  4492  	}
  4493  }
  4494  
  4495  func TestStateStore_UpdateAllocsFromClient(t *testing.T) {
  4496  	t.Parallel()
  4497  
  4498  	state := testStateStore(t)
  4499  	parent := mock.Job()
  4500  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 998, parent); err != nil {
  4501  		t.Fatalf("err: %v", err)
  4502  	}
  4503  
  4504  	child := mock.Job()
  4505  	child.Status = ""
  4506  	child.ParentID = parent.ID
  4507  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, child); err != nil {
  4508  		t.Fatalf("err: %v", err)
  4509  	}
  4510  
  4511  	alloc := mock.Alloc()
  4512  	alloc.JobID = child.ID
  4513  	alloc.Job = child
  4514  
  4515  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  4516  	if err != nil {
  4517  		t.Fatalf("err: %v", err)
  4518  	}
  4519  
  4520  	ws := memdb.NewWatchSet()
  4521  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  4522  	if err != nil {
  4523  		t.Fatalf("err: %v", err)
  4524  	}
  4525  	if summary == nil {
  4526  		t.Fatalf("nil summary")
  4527  	}
  4528  	if summary.JobID != parent.ID {
  4529  		t.Fatalf("bad summary id: %v", parent.ID)
  4530  	}
  4531  	if summary.Children == nil {
  4532  		t.Fatalf("nil children summary")
  4533  	}
  4534  	if summary.Children.Pending != 0 || summary.Children.Running != 1 || summary.Children.Dead != 0 {
  4535  		t.Fatalf("bad children summary: %v", summary.Children)
  4536  	}
  4537  
  4538  	// Create watchsets so we can test that update fires the watch
  4539  	ws = memdb.NewWatchSet()
  4540  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  4541  		t.Fatalf("bad: %v", err)
  4542  	}
  4543  
  4544  	// Create the delta updates
  4545  	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStateRunning}}
  4546  	update := &structs.Allocation{
  4547  		ID:           alloc.ID,
  4548  		ClientStatus: structs.AllocClientStatusComplete,
  4549  		TaskStates:   ts,
  4550  		JobID:        alloc.JobID,
  4551  		TaskGroup:    alloc.TaskGroup,
  4552  	}
  4553  	err = state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update})
  4554  	if err != nil {
  4555  		t.Fatalf("err: %v", err)
  4556  	}
  4557  
  4558  	if !watchFired(ws) {
  4559  		t.Fatalf("bad")
  4560  	}
  4561  
  4562  	ws = memdb.NewWatchSet()
  4563  	summary, err = state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  4564  	if err != nil {
  4565  		t.Fatalf("err: %v", err)
  4566  	}
  4567  	if summary == nil {
  4568  		t.Fatalf("nil summary")
  4569  	}
  4570  	if summary.JobID != parent.ID {
  4571  		t.Fatalf("bad summary id: %v", parent.ID)
  4572  	}
  4573  	if summary.Children == nil {
  4574  		t.Fatalf("nil children summary")
  4575  	}
  4576  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  4577  		t.Fatalf("bad children summary: %v", summary.Children)
  4578  	}
  4579  
  4580  	if watchFired(ws) {
  4581  		t.Fatalf("bad")
  4582  	}
  4583  }
  4584  
  4585  func TestStateStore_UpdateAllocsFromClient_ChildJob(t *testing.T) {
  4586  	t.Parallel()
  4587  
  4588  	state := testStateStore(t)
  4589  	alloc1 := mock.Alloc()
  4590  	alloc2 := mock.Alloc()
  4591  
  4592  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc1.Job); err != nil {
  4593  		t.Fatalf("err: %v", err)
  4594  	}
  4595  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc2.Job); err != nil {
  4596  		t.Fatalf("err: %v", err)
  4597  	}
  4598  
  4599  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc1, alloc2})
  4600  	if err != nil {
  4601  		t.Fatalf("err: %v", err)
  4602  	}
  4603  
  4604  	// Create watchsets so we can test that update fires the watch
  4605  	watches := make([]memdb.WatchSet, 8)
  4606  	for i := 0; i < 8; i++ {
  4607  		watches[i] = memdb.NewWatchSet()
  4608  	}
  4609  	if _, err := state.AllocByID(watches[0], alloc1.ID); err != nil {
  4610  		t.Fatalf("bad: %v", err)
  4611  	}
  4612  	if _, err := state.AllocByID(watches[1], alloc2.ID); err != nil {
  4613  		t.Fatalf("bad: %v", err)
  4614  	}
  4615  	if _, err := state.AllocsByEval(watches[2], alloc1.EvalID); err != nil {
  4616  		t.Fatalf("bad: %v", err)
  4617  	}
  4618  	if _, err := state.AllocsByEval(watches[3], alloc2.EvalID); err != nil {
  4619  		t.Fatalf("bad: %v", err)
  4620  	}
  4621  	if _, err := state.AllocsByJob(watches[4], alloc1.Namespace, alloc1.JobID, false); err != nil {
  4622  		t.Fatalf("bad: %v", err)
  4623  	}
  4624  	if _, err := state.AllocsByJob(watches[5], alloc2.Namespace, alloc2.JobID, false); err != nil {
  4625  		t.Fatalf("bad: %v", err)
  4626  	}
  4627  	if _, err := state.AllocsByNode(watches[6], alloc1.NodeID); err != nil {
  4628  		t.Fatalf("bad: %v", err)
  4629  	}
  4630  	if _, err := state.AllocsByNode(watches[7], alloc2.NodeID); err != nil {
  4631  		t.Fatalf("bad: %v", err)
  4632  	}
  4633  
  4634  	// Create the delta updates
  4635  	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStatePending}}
  4636  	update := &structs.Allocation{
  4637  		ID:           alloc1.ID,
  4638  		ClientStatus: structs.AllocClientStatusFailed,
  4639  		TaskStates:   ts,
  4640  		JobID:        alloc1.JobID,
  4641  		TaskGroup:    alloc1.TaskGroup,
  4642  	}
  4643  	update2 := &structs.Allocation{
  4644  		ID:           alloc2.ID,
  4645  		ClientStatus: structs.AllocClientStatusRunning,
  4646  		TaskStates:   ts,
  4647  		JobID:        alloc2.JobID,
  4648  		TaskGroup:    alloc2.TaskGroup,
  4649  	}
  4650  
  4651  	err = state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update, update2})
  4652  	if err != nil {
  4653  		t.Fatalf("err: %v", err)
  4654  	}
  4655  
  4656  	for i, ws := range watches {
  4657  		if !watchFired(ws) {
  4658  			t.Fatalf("bad %d", i)
  4659  		}
  4660  	}
  4661  
  4662  	ws := memdb.NewWatchSet()
  4663  	out, err := state.AllocByID(ws, alloc1.ID)
  4664  	if err != nil {
  4665  		t.Fatalf("err: %v", err)
  4666  	}
  4667  
  4668  	alloc1.CreateIndex = 1000
  4669  	alloc1.ModifyIndex = 1001
  4670  	alloc1.TaskStates = ts
  4671  	alloc1.ClientStatus = structs.AllocClientStatusFailed
  4672  	if !reflect.DeepEqual(alloc1, out) {
  4673  		t.Fatalf("bad: %#v %#v", alloc1, out)
  4674  	}
  4675  
  4676  	out, err = state.AllocByID(ws, alloc2.ID)
  4677  	if err != nil {
  4678  		t.Fatalf("err: %v", err)
  4679  	}
  4680  
  4681  	alloc2.ModifyIndex = 1000
  4682  	alloc2.ModifyIndex = 1001
  4683  	alloc2.ClientStatus = structs.AllocClientStatusRunning
  4684  	alloc2.TaskStates = ts
  4685  	if !reflect.DeepEqual(alloc2, out) {
  4686  		t.Fatalf("bad: %#v %#v", alloc2, out)
  4687  	}
  4688  
  4689  	index, err := state.Index("allocs")
  4690  	if err != nil {
  4691  		t.Fatalf("err: %v", err)
  4692  	}
  4693  	if index != 1001 {
  4694  		t.Fatalf("bad: %d", index)
  4695  	}
  4696  
  4697  	// Ensure summaries have been updated
  4698  	summary, err := state.JobSummaryByID(ws, alloc1.Namespace, alloc1.JobID)
  4699  	if err != nil {
  4700  		t.Fatalf("err: %v", err)
  4701  	}
  4702  	tgSummary := summary.Summary["web"]
  4703  	if tgSummary.Failed != 1 {
  4704  		t.Fatalf("expected failed: %v, actual: %v, summary: %#v", 1, tgSummary.Failed, tgSummary)
  4705  	}
  4706  
  4707  	summary2, err := state.JobSummaryByID(ws, alloc2.Namespace, alloc2.JobID)
  4708  	if err != nil {
  4709  		t.Fatalf("err: %v", err)
  4710  	}
  4711  	tgSummary2 := summary2.Summary["web"]
  4712  	if tgSummary2.Running != 1 {
  4713  		t.Fatalf("expected running: %v, actual: %v", 1, tgSummary2.Running)
  4714  	}
  4715  
  4716  	if watchFired(ws) {
  4717  		t.Fatalf("bad")
  4718  	}
  4719  }
  4720  
  4721  func TestStateStore_UpdateMultipleAllocsFromClient(t *testing.T) {
  4722  	t.Parallel()
  4723  
  4724  	state := testStateStore(t)
  4725  	alloc := mock.Alloc()
  4726  
  4727  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
  4728  		t.Fatalf("err: %v", err)
  4729  	}
  4730  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  4731  	if err != nil {
  4732  		t.Fatalf("err: %v", err)
  4733  	}
  4734  
  4735  	// Create the delta updates
  4736  	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStatePending}}
  4737  	update := &structs.Allocation{
  4738  		ID:           alloc.ID,
  4739  		ClientStatus: structs.AllocClientStatusRunning,
  4740  		TaskStates:   ts,
  4741  		JobID:        alloc.JobID,
  4742  		TaskGroup:    alloc.TaskGroup,
  4743  	}
  4744  	update2 := &structs.Allocation{
  4745  		ID:           alloc.ID,
  4746  		ClientStatus: structs.AllocClientStatusPending,
  4747  		TaskStates:   ts,
  4748  		JobID:        alloc.JobID,
  4749  		TaskGroup:    alloc.TaskGroup,
  4750  	}
  4751  
  4752  	err = state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update, update2})
  4753  	if err != nil {
  4754  		t.Fatalf("err: %v", err)
  4755  	}
  4756  
  4757  	ws := memdb.NewWatchSet()
  4758  	out, err := state.AllocByID(ws, alloc.ID)
  4759  	if err != nil {
  4760  		t.Fatalf("err: %v", err)
  4761  	}
  4762  
  4763  	alloc.CreateIndex = 1000
  4764  	alloc.ModifyIndex = 1001
  4765  	alloc.TaskStates = ts
  4766  	alloc.ClientStatus = structs.AllocClientStatusPending
  4767  	if !reflect.DeepEqual(alloc, out) {
  4768  		t.Fatalf("bad: %#v , actual:%#v", alloc, out)
  4769  	}
  4770  
  4771  	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  4772  	expectedSummary := &structs.JobSummary{
  4773  		JobID:     alloc.JobID,
  4774  		Namespace: alloc.Namespace,
  4775  		Summary: map[string]structs.TaskGroupSummary{
  4776  			"web": {
  4777  				Starting: 1,
  4778  			},
  4779  		},
  4780  		Children:    new(structs.JobChildrenSummary),
  4781  		CreateIndex: 999,
  4782  		ModifyIndex: 1001,
  4783  	}
  4784  	if err != nil {
  4785  		t.Fatalf("err: %v", err)
  4786  	}
  4787  	if !reflect.DeepEqual(summary, expectedSummary) {
  4788  		t.Fatalf("expected: %#v, actual: %#v", expectedSummary, summary)
  4789  	}
  4790  }
  4791  
  4792  func TestStateStore_UpdateAllocsFromClient_Deployment(t *testing.T) {
  4793  	t.Parallel()
  4794  	require := require.New(t)
  4795  
  4796  	state := testStateStore(t)
  4797  
  4798  	alloc := mock.Alloc()
  4799  	now := time.Now()
  4800  	alloc.CreateTime = now.UnixNano()
  4801  	pdeadline := 5 * time.Minute
  4802  	deployment := mock.Deployment()
  4803  	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
  4804  	alloc.DeploymentID = deployment.ID
  4805  
  4806  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
  4807  	require.Nil(state.UpsertDeployment(1000, deployment))
  4808  	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc}))
  4809  
  4810  	healthy := now.Add(time.Second)
  4811  	update := &structs.Allocation{
  4812  		ID:           alloc.ID,
  4813  		ClientStatus: structs.AllocClientStatusRunning,
  4814  		JobID:        alloc.JobID,
  4815  		TaskGroup:    alloc.TaskGroup,
  4816  		DeploymentStatus: &structs.AllocDeploymentStatus{
  4817  			Healthy:   helper.BoolToPtr(true),
  4818  			Timestamp: healthy,
  4819  		},
  4820  	}
  4821  	require.Nil(state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update}))
  4822  
  4823  	// Check that the deployment state was updated because the healthy
  4824  	// deployment
  4825  	dout, err := state.DeploymentByID(nil, deployment.ID)
  4826  	require.Nil(err)
  4827  	require.NotNil(dout)
  4828  	require.Len(dout.TaskGroups, 1)
  4829  	dstate := dout.TaskGroups[alloc.TaskGroup]
  4830  	require.NotNil(dstate)
  4831  	require.Equal(1, dstate.PlacedAllocs)
  4832  	require.True(healthy.Add(pdeadline).Equal(dstate.RequireProgressBy))
  4833  }
  4834  
  4835  // This tests that the deployment state is merged correctly
  4836  func TestStateStore_UpdateAllocsFromClient_DeploymentStateMerges(t *testing.T) {
  4837  	t.Parallel()
  4838  	require := require.New(t)
  4839  
  4840  	state := testStateStore(t)
  4841  	alloc := mock.Alloc()
  4842  	now := time.Now()
  4843  	alloc.CreateTime = now.UnixNano()
  4844  	pdeadline := 5 * time.Minute
  4845  	deployment := mock.Deployment()
  4846  	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
  4847  	alloc.DeploymentID = deployment.ID
  4848  	alloc.DeploymentStatus = &structs.AllocDeploymentStatus{
  4849  		Canary: true,
  4850  	}
  4851  
  4852  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
  4853  	require.Nil(state.UpsertDeployment(1000, deployment))
  4854  	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc}))
  4855  
  4856  	update := &structs.Allocation{
  4857  		ID:           alloc.ID,
  4858  		ClientStatus: structs.AllocClientStatusRunning,
  4859  		JobID:        alloc.JobID,
  4860  		TaskGroup:    alloc.TaskGroup,
  4861  		DeploymentStatus: &structs.AllocDeploymentStatus{
  4862  			Healthy: helper.BoolToPtr(true),
  4863  			Canary:  false,
  4864  		},
  4865  	}
  4866  	require.Nil(state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{update}))
  4867  
  4868  	// Check that the merging of the deployment status was correct
  4869  	out, err := state.AllocByID(nil, alloc.ID)
  4870  	require.Nil(err)
  4871  	require.NotNil(out)
  4872  	require.True(out.DeploymentStatus.Canary)
  4873  	require.NotNil(out.DeploymentStatus.Healthy)
  4874  	require.True(*out.DeploymentStatus.Healthy)
  4875  }
  4876  
  4877  func TestStateStore_UpsertAlloc_Alloc(t *testing.T) {
  4878  	t.Parallel()
  4879  
  4880  	state := testStateStore(t)
  4881  	alloc := mock.Alloc()
  4882  
  4883  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
  4884  		t.Fatalf("err: %v", err)
  4885  	}
  4886  
  4887  	// Create watchsets so we can test that update fires the watch
  4888  	watches := make([]memdb.WatchSet, 4)
  4889  	for i := 0; i < 4; i++ {
  4890  		watches[i] = memdb.NewWatchSet()
  4891  	}
  4892  	if _, err := state.AllocByID(watches[0], alloc.ID); err != nil {
  4893  		t.Fatalf("bad: %v", err)
  4894  	}
  4895  	if _, err := state.AllocsByEval(watches[1], alloc.EvalID); err != nil {
  4896  		t.Fatalf("bad: %v", err)
  4897  	}
  4898  	if _, err := state.AllocsByJob(watches[2], alloc.Namespace, alloc.JobID, false); err != nil {
  4899  		t.Fatalf("bad: %v", err)
  4900  	}
  4901  	if _, err := state.AllocsByNode(watches[3], alloc.NodeID); err != nil {
  4902  		t.Fatalf("bad: %v", err)
  4903  	}
  4904  
  4905  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  4906  	if err != nil {
  4907  		t.Fatalf("err: %v", err)
  4908  	}
  4909  
  4910  	for i, ws := range watches {
  4911  		if !watchFired(ws) {
  4912  			t.Fatalf("bad %d", i)
  4913  		}
  4914  	}
  4915  
  4916  	ws := memdb.NewWatchSet()
  4917  	out, err := state.AllocByID(ws, alloc.ID)
  4918  	if err != nil {
  4919  		t.Fatalf("err: %v", err)
  4920  	}
  4921  
  4922  	if !reflect.DeepEqual(alloc, out) {
  4923  		t.Fatalf("bad: %#v %#v", alloc, out)
  4924  	}
  4925  
  4926  	index, err := state.Index("allocs")
  4927  	if err != nil {
  4928  		t.Fatalf("err: %v", err)
  4929  	}
  4930  	if index != 1000 {
  4931  		t.Fatalf("bad: %d", index)
  4932  	}
  4933  
  4934  	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  4935  	if err != nil {
  4936  		t.Fatalf("err: %v", err)
  4937  	}
  4938  
  4939  	tgSummary, ok := summary.Summary["web"]
  4940  	if !ok {
  4941  		t.Fatalf("no summary for task group web")
  4942  	}
  4943  	if tgSummary.Starting != 1 {
  4944  		t.Fatalf("expected queued: %v, actual: %v", 1, tgSummary.Starting)
  4945  	}
  4946  
  4947  	if watchFired(ws) {
  4948  		t.Fatalf("bad")
  4949  	}
  4950  }
  4951  
  4952  func TestStateStore_UpsertAlloc_Deployment(t *testing.T) {
  4953  	t.Parallel()
  4954  	require := require.New(t)
  4955  
  4956  	state := testStateStore(t)
  4957  	alloc := mock.Alloc()
  4958  	now := time.Now()
  4959  	alloc.CreateTime = now.UnixNano()
  4960  	alloc.ModifyTime = now.UnixNano()
  4961  	pdeadline := 5 * time.Minute
  4962  	deployment := mock.Deployment()
  4963  	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
  4964  	alloc.DeploymentID = deployment.ID
  4965  
  4966  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
  4967  	require.Nil(state.UpsertDeployment(1000, deployment))
  4968  
  4969  	// Create a watch set so we can test that update fires the watch
  4970  	ws := memdb.NewWatchSet()
  4971  	require.Nil(state.AllocsByDeployment(ws, alloc.DeploymentID))
  4972  
  4973  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc})
  4974  	require.Nil(err)
  4975  
  4976  	if !watchFired(ws) {
  4977  		t.Fatalf("watch not fired")
  4978  	}
  4979  
  4980  	ws = memdb.NewWatchSet()
  4981  	allocs, err := state.AllocsByDeployment(ws, alloc.DeploymentID)
  4982  	require.Nil(err)
  4983  	require.Len(allocs, 1)
  4984  	require.EqualValues(alloc, allocs[0])
  4985  
  4986  	index, err := state.Index("allocs")
  4987  	require.Nil(err)
  4988  	require.EqualValues(1001, index)
  4989  	if watchFired(ws) {
  4990  		t.Fatalf("bad")
  4991  	}
  4992  
  4993  	// Check that the deployment state was updated
  4994  	dout, err := state.DeploymentByID(nil, deployment.ID)
  4995  	require.Nil(err)
  4996  	require.NotNil(dout)
  4997  	require.Len(dout.TaskGroups, 1)
  4998  	dstate := dout.TaskGroups[alloc.TaskGroup]
  4999  	require.NotNil(dstate)
  5000  	require.Equal(1, dstate.PlacedAllocs)
  5001  	require.True(now.Add(pdeadline).Equal(dstate.RequireProgressBy))
  5002  }
  5003  
  5004  // Testing to ensure we keep issue
  5005  // https://github.com/hashicorp/nomad/issues/2583 fixed
  5006  func TestStateStore_UpsertAlloc_No_Job(t *testing.T) {
  5007  	t.Parallel()
  5008  
  5009  	state := testStateStore(t)
  5010  	alloc := mock.Alloc()
  5011  	alloc.Job = nil
  5012  
  5013  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 999, []*structs.Allocation{alloc})
  5014  	if err == nil || !strings.Contains(err.Error(), "without a job") {
  5015  		t.Fatalf("expect err: %v", err)
  5016  	}
  5017  }
  5018  
  5019  func TestStateStore_UpsertAlloc_ChildJob(t *testing.T) {
  5020  	t.Parallel()
  5021  
  5022  	state := testStateStore(t)
  5023  
  5024  	parent := mock.Job()
  5025  	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 998, parent))
  5026  
  5027  	child := mock.Job()
  5028  	child.Status = ""
  5029  	child.ParentID = parent.ID
  5030  
  5031  	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 999, child))
  5032  
  5033  	alloc := mock.Alloc()
  5034  	alloc.JobID = child.ID
  5035  	alloc.Job = child
  5036  
  5037  	// Create watchsets so we can test that delete fires the watch
  5038  	ws := memdb.NewWatchSet()
  5039  	_, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  5040  	require.NoError(t, err)
  5041  
  5042  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  5043  	require.NoError(t, err)
  5044  
  5045  	require.True(t, watchFired(ws))
  5046  
  5047  	ws = memdb.NewWatchSet()
  5048  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  5049  	require.NoError(t, err)
  5050  	require.NotNil(t, summary)
  5051  
  5052  	require.Equal(t, parent.ID, summary.JobID)
  5053  	require.NotNil(t, summary.Children)
  5054  
  5055  	require.Equal(t, int64(0), summary.Children.Pending)
  5056  	require.Equal(t, int64(1), summary.Children.Running)
  5057  	require.Equal(t, int64(0), summary.Children.Dead)
  5058  
  5059  	require.False(t, watchFired(ws))
  5060  }
  5061  
  5062  func TestStateStore_UpdateAlloc_Alloc(t *testing.T) {
  5063  	t.Parallel()
  5064  
  5065  	state := testStateStore(t)
  5066  	alloc := mock.Alloc()
  5067  
  5068  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
  5069  		t.Fatalf("err: %v", err)
  5070  	}
  5071  
  5072  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  5073  	if err != nil {
  5074  		t.Fatalf("err: %v", err)
  5075  	}
  5076  
  5077  	ws := memdb.NewWatchSet()
  5078  	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  5079  	if err != nil {
  5080  		t.Fatalf("err: %v", err)
  5081  	}
  5082  	tgSummary := summary.Summary["web"]
  5083  	if tgSummary.Starting != 1 {
  5084  		t.Fatalf("expected starting: %v, actual: %v", 1, tgSummary.Starting)
  5085  	}
  5086  
  5087  	alloc2 := mock.Alloc()
  5088  	alloc2.ID = alloc.ID
  5089  	alloc2.NodeID = alloc.NodeID + ".new"
  5090  	state.UpsertJobSummary(1001, mock.JobSummary(alloc2.JobID))
  5091  
  5092  	// Create watchsets so we can test that update fires the watch
  5093  	watches := make([]memdb.WatchSet, 4)
  5094  	for i := 0; i < 4; i++ {
  5095  		watches[i] = memdb.NewWatchSet()
  5096  	}
  5097  	if _, err := state.AllocByID(watches[0], alloc2.ID); err != nil {
  5098  		t.Fatalf("bad: %v", err)
  5099  	}
  5100  	if _, err := state.AllocsByEval(watches[1], alloc2.EvalID); err != nil {
  5101  		t.Fatalf("bad: %v", err)
  5102  	}
  5103  	if _, err := state.AllocsByJob(watches[2], alloc2.Namespace, alloc2.JobID, false); err != nil {
  5104  		t.Fatalf("bad: %v", err)
  5105  	}
  5106  	if _, err := state.AllocsByNode(watches[3], alloc2.NodeID); err != nil {
  5107  		t.Fatalf("bad: %v", err)
  5108  	}
  5109  
  5110  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{alloc2})
  5111  	if err != nil {
  5112  		t.Fatalf("err: %v", err)
  5113  	}
  5114  
  5115  	for i, ws := range watches {
  5116  		if !watchFired(ws) {
  5117  			t.Fatalf("bad %d", i)
  5118  		}
  5119  	}
  5120  
  5121  	ws = memdb.NewWatchSet()
  5122  	out, err := state.AllocByID(ws, alloc.ID)
  5123  	if err != nil {
  5124  		t.Fatalf("err: %v", err)
  5125  	}
  5126  
  5127  	if !reflect.DeepEqual(alloc2, out) {
  5128  		t.Fatalf("bad: %#v %#v", alloc2, out)
  5129  	}
  5130  
  5131  	if out.CreateIndex != 1000 {
  5132  		t.Fatalf("bad: %#v", out)
  5133  	}
  5134  	if out.ModifyIndex != 1002 {
  5135  		t.Fatalf("bad: %#v", out)
  5136  	}
  5137  
  5138  	index, err := state.Index("allocs")
  5139  	if err != nil {
  5140  		t.Fatalf("err: %v", err)
  5141  	}
  5142  	if index != 1002 {
  5143  		t.Fatalf("bad: %d", index)
  5144  	}
  5145  
  5146  	// Ensure that summary hasb't changed
  5147  	summary, err = state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  5148  	if err != nil {
  5149  		t.Fatalf("err: %v", err)
  5150  	}
  5151  	tgSummary = summary.Summary["web"]
  5152  	if tgSummary.Starting != 1 {
  5153  		t.Fatalf("expected starting: %v, actual: %v", 1, tgSummary.Starting)
  5154  	}
  5155  
  5156  	if watchFired(ws) {
  5157  		t.Fatalf("bad")
  5158  	}
  5159  }
  5160  
  5161  // This test ensures that the state store will mark the clients status as lost
  5162  // when set rather than preferring the existing status.
  5163  func TestStateStore_UpdateAlloc_Lost(t *testing.T) {
  5164  	t.Parallel()
  5165  
  5166  	state := testStateStore(t)
  5167  	alloc := mock.Alloc()
  5168  	alloc.ClientStatus = "foo"
  5169  
  5170  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
  5171  		t.Fatalf("err: %v", err)
  5172  	}
  5173  
  5174  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  5175  	if err != nil {
  5176  		t.Fatalf("err: %v", err)
  5177  	}
  5178  
  5179  	alloc2 := new(structs.Allocation)
  5180  	*alloc2 = *alloc
  5181  	alloc2.ClientStatus = structs.AllocClientStatusLost
  5182  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc2}); err != nil {
  5183  		t.Fatalf("err: %v", err)
  5184  	}
  5185  
  5186  	ws := memdb.NewWatchSet()
  5187  	out, err := state.AllocByID(ws, alloc2.ID)
  5188  	if err != nil {
  5189  		t.Fatalf("err: %v", err)
  5190  	}
  5191  
  5192  	if out.ClientStatus != structs.AllocClientStatusLost {
  5193  		t.Fatalf("bad: %#v", out)
  5194  	}
  5195  }
  5196  
  5197  // This test ensures an allocation can be updated when there is no job
  5198  // associated with it. This will happen when a job is stopped by an user which
  5199  // has non-terminal allocations on clients
  5200  func TestStateStore_UpdateAlloc_NoJob(t *testing.T) {
  5201  	t.Parallel()
  5202  
  5203  	state := testStateStore(t)
  5204  	alloc := mock.Alloc()
  5205  
  5206  	// Upsert a job
  5207  	state.UpsertJobSummary(998, mock.JobSummary(alloc.JobID))
  5208  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job); err != nil {
  5209  		t.Fatalf("err: %v", err)
  5210  	}
  5211  
  5212  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  5213  	if err != nil {
  5214  		t.Fatalf("err: %v", err)
  5215  	}
  5216  
  5217  	if err := state.DeleteJob(1001, alloc.Namespace, alloc.JobID); err != nil {
  5218  		t.Fatalf("err: %v", err)
  5219  	}
  5220  
  5221  	// Update the desired state of the allocation to stop
  5222  	allocCopy := alloc.Copy()
  5223  	allocCopy.DesiredStatus = structs.AllocDesiredStatusStop
  5224  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{allocCopy}); err != nil {
  5225  		t.Fatalf("err: %v", err)
  5226  	}
  5227  
  5228  	// Update the client state of the allocation to complete
  5229  	allocCopy1 := allocCopy.Copy()
  5230  	allocCopy1.ClientStatus = structs.AllocClientStatusComplete
  5231  	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1003, []*structs.Allocation{allocCopy1}); err != nil {
  5232  		t.Fatalf("err: %v", err)
  5233  	}
  5234  
  5235  	ws := memdb.NewWatchSet()
  5236  	out, _ := state.AllocByID(ws, alloc.ID)
  5237  	// Update the modify index of the alloc before comparing
  5238  	allocCopy1.ModifyIndex = 1003
  5239  	if !reflect.DeepEqual(out, allocCopy1) {
  5240  		t.Fatalf("expected: %#v \n actual: %#v", allocCopy1, out)
  5241  	}
  5242  }
  5243  
  5244  func TestStateStore_UpdateAllocDesiredTransition(t *testing.T) {
  5245  	t.Parallel()
  5246  	require := require.New(t)
  5247  
  5248  	state := testStateStore(t)
  5249  	alloc := mock.Alloc()
  5250  
  5251  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job))
  5252  	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc}))
  5253  
  5254  	t1 := &structs.DesiredTransition{
  5255  		Migrate: helper.BoolToPtr(true),
  5256  	}
  5257  	t2 := &structs.DesiredTransition{
  5258  		Migrate: helper.BoolToPtr(false),
  5259  	}
  5260  	eval := &structs.Evaluation{
  5261  		ID:             uuid.Generate(),
  5262  		Namespace:      alloc.Namespace,
  5263  		Priority:       alloc.Job.Priority,
  5264  		Type:           alloc.Job.Type,
  5265  		TriggeredBy:    structs.EvalTriggerNodeDrain,
  5266  		JobID:          alloc.Job.ID,
  5267  		JobModifyIndex: alloc.Job.ModifyIndex,
  5268  		Status:         structs.EvalStatusPending,
  5269  	}
  5270  	evals := []*structs.Evaluation{eval}
  5271  
  5272  	m := map[string]*structs.DesiredTransition{alloc.ID: t1}
  5273  	require.Nil(state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, 1001, m, evals))
  5274  
  5275  	ws := memdb.NewWatchSet()
  5276  	out, err := state.AllocByID(ws, alloc.ID)
  5277  	require.Nil(err)
  5278  	require.NotNil(out.DesiredTransition.Migrate)
  5279  	require.True(*out.DesiredTransition.Migrate)
  5280  	require.EqualValues(1000, out.CreateIndex)
  5281  	require.EqualValues(1001, out.ModifyIndex)
  5282  
  5283  	index, err := state.Index("allocs")
  5284  	require.Nil(err)
  5285  	require.EqualValues(1001, index)
  5286  
  5287  	// Check the eval is created
  5288  	eout, err := state.EvalByID(nil, eval.ID)
  5289  	require.Nil(err)
  5290  	require.NotNil(eout)
  5291  
  5292  	m = map[string]*structs.DesiredTransition{alloc.ID: t2}
  5293  	require.Nil(state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, 1002, m, evals))
  5294  
  5295  	ws = memdb.NewWatchSet()
  5296  	out, err = state.AllocByID(ws, alloc.ID)
  5297  	require.Nil(err)
  5298  	require.NotNil(out.DesiredTransition.Migrate)
  5299  	require.False(*out.DesiredTransition.Migrate)
  5300  	require.EqualValues(1000, out.CreateIndex)
  5301  	require.EqualValues(1002, out.ModifyIndex)
  5302  
  5303  	index, err = state.Index("allocs")
  5304  	require.Nil(err)
  5305  	require.EqualValues(1002, index)
  5306  
  5307  	// Try with a bogus alloc id
  5308  	m = map[string]*structs.DesiredTransition{uuid.Generate(): t2}
  5309  	require.Nil(state.UpdateAllocsDesiredTransitions(structs.MsgTypeTestSetup, 1003, m, evals))
  5310  }
  5311  
  5312  func TestStateStore_JobSummary(t *testing.T) {
  5313  	t.Parallel()
  5314  
  5315  	state := testStateStore(t)
  5316  
  5317  	// Add a job
  5318  	job := mock.Job()
  5319  	state.UpsertJob(structs.MsgTypeTestSetup, 900, job)
  5320  
  5321  	// Get the job back
  5322  	ws := memdb.NewWatchSet()
  5323  	outJob, _ := state.JobByID(ws, job.Namespace, job.ID)
  5324  	if outJob.CreateIndex != 900 {
  5325  		t.Fatalf("bad create index: %v", outJob.CreateIndex)
  5326  	}
  5327  	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
  5328  	if summary.CreateIndex != 900 {
  5329  		t.Fatalf("bad create index: %v", summary.CreateIndex)
  5330  	}
  5331  
  5332  	// Upsert an allocation
  5333  	alloc := mock.Alloc()
  5334  	alloc.JobID = job.ID
  5335  	alloc.Job = job
  5336  	state.UpsertAllocs(structs.MsgTypeTestSetup, 910, []*structs.Allocation{alloc})
  5337  
  5338  	// Update the alloc from client
  5339  	alloc1 := alloc.Copy()
  5340  	alloc1.ClientStatus = structs.AllocClientStatusPending
  5341  	alloc1.DesiredStatus = ""
  5342  	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 920, []*structs.Allocation{alloc})
  5343  
  5344  	alloc3 := alloc.Copy()
  5345  	alloc3.ClientStatus = structs.AllocClientStatusRunning
  5346  	alloc3.DesiredStatus = ""
  5347  	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 930, []*structs.Allocation{alloc3})
  5348  
  5349  	// Upsert the alloc
  5350  	alloc4 := alloc.Copy()
  5351  	alloc4.ClientStatus = structs.AllocClientStatusPending
  5352  	alloc4.DesiredStatus = structs.AllocDesiredStatusRun
  5353  	state.UpsertAllocs(structs.MsgTypeTestSetup, 950, []*structs.Allocation{alloc4})
  5354  
  5355  	// Again upsert the alloc
  5356  	alloc5 := alloc.Copy()
  5357  	alloc5.ClientStatus = structs.AllocClientStatusPending
  5358  	alloc5.DesiredStatus = structs.AllocDesiredStatusRun
  5359  	state.UpsertAllocs(structs.MsgTypeTestSetup, 970, []*structs.Allocation{alloc5})
  5360  
  5361  	if !watchFired(ws) {
  5362  		t.Fatalf("bad")
  5363  	}
  5364  
  5365  	expectedSummary := structs.JobSummary{
  5366  		JobID:     job.ID,
  5367  		Namespace: job.Namespace,
  5368  		Summary: map[string]structs.TaskGroupSummary{
  5369  			"web": {
  5370  				Running: 1,
  5371  			},
  5372  		},
  5373  		Children:    new(structs.JobChildrenSummary),
  5374  		CreateIndex: 900,
  5375  		ModifyIndex: 930,
  5376  	}
  5377  
  5378  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  5379  	if !reflect.DeepEqual(&expectedSummary, summary) {
  5380  		t.Fatalf("expected: %#v, actual: %v", expectedSummary, summary)
  5381  	}
  5382  
  5383  	// De-register the job.
  5384  	state.DeleteJob(980, job.Namespace, job.ID)
  5385  
  5386  	// Shouldn't have any effect on the summary
  5387  	alloc6 := alloc.Copy()
  5388  	alloc6.ClientStatus = structs.AllocClientStatusRunning
  5389  	alloc6.DesiredStatus = ""
  5390  	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 990, []*structs.Allocation{alloc6})
  5391  
  5392  	// We shouldn't have any summary at this point
  5393  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  5394  	if summary != nil {
  5395  		t.Fatalf("expected nil, actual: %#v", summary)
  5396  	}
  5397  
  5398  	// Re-register the same job
  5399  	job1 := mock.Job()
  5400  	job1.ID = job.ID
  5401  	state.UpsertJob(structs.MsgTypeTestSetup, 1000, job1)
  5402  	outJob2, _ := state.JobByID(ws, job1.Namespace, job1.ID)
  5403  	if outJob2.CreateIndex != 1000 {
  5404  		t.Fatalf("bad create index: %v", outJob2.CreateIndex)
  5405  	}
  5406  	summary, _ = state.JobSummaryByID(ws, job1.Namespace, job1.ID)
  5407  	if summary.CreateIndex != 1000 {
  5408  		t.Fatalf("bad create index: %v", summary.CreateIndex)
  5409  	}
  5410  
  5411  	// Upsert an allocation
  5412  	alloc7 := alloc.Copy()
  5413  	alloc7.JobID = outJob.ID
  5414  	alloc7.Job = outJob
  5415  	alloc7.ClientStatus = structs.AllocClientStatusComplete
  5416  	alloc7.DesiredStatus = structs.AllocDesiredStatusRun
  5417  	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1020, []*structs.Allocation{alloc7})
  5418  
  5419  	expectedSummary = structs.JobSummary{
  5420  		JobID:     job.ID,
  5421  		Namespace: job.Namespace,
  5422  		Summary: map[string]structs.TaskGroupSummary{
  5423  			"web": {},
  5424  		},
  5425  		Children:    new(structs.JobChildrenSummary),
  5426  		CreateIndex: 1000,
  5427  		ModifyIndex: 1000,
  5428  	}
  5429  
  5430  	summary, _ = state.JobSummaryByID(ws, job1.Namespace, job1.ID)
  5431  	if !reflect.DeepEqual(&expectedSummary, summary) {
  5432  		t.Fatalf("expected: %#v, actual: %#v", expectedSummary, summary)
  5433  	}
  5434  }
  5435  
  5436  func TestStateStore_ReconcileJobSummary(t *testing.T) {
  5437  	t.Parallel()
  5438  
  5439  	state := testStateStore(t)
  5440  
  5441  	// Create an alloc
  5442  	alloc := mock.Alloc()
  5443  
  5444  	// Add another task group to the job
  5445  	tg2 := alloc.Job.TaskGroups[0].Copy()
  5446  	tg2.Name = "db"
  5447  	alloc.Job.TaskGroups = append(alloc.Job.TaskGroups, tg2)
  5448  	state.UpsertJob(structs.MsgTypeTestSetup, 100, alloc.Job)
  5449  
  5450  	// Create one more alloc for the db task group
  5451  	alloc2 := mock.Alloc()
  5452  	alloc2.TaskGroup = "db"
  5453  	alloc2.JobID = alloc.JobID
  5454  	alloc2.Job = alloc.Job
  5455  
  5456  	// Upserts the alloc
  5457  	state.UpsertAllocs(structs.MsgTypeTestSetup, 110, []*structs.Allocation{alloc, alloc2})
  5458  
  5459  	// Change the state of the first alloc to running
  5460  	alloc3 := alloc.Copy()
  5461  	alloc3.ClientStatus = structs.AllocClientStatusRunning
  5462  	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 120, []*structs.Allocation{alloc3})
  5463  
  5464  	//Add some more allocs to the second tg
  5465  	alloc4 := mock.Alloc()
  5466  	alloc4.JobID = alloc.JobID
  5467  	alloc4.Job = alloc.Job
  5468  	alloc4.TaskGroup = "db"
  5469  	alloc5 := alloc4.Copy()
  5470  	alloc5.ClientStatus = structs.AllocClientStatusRunning
  5471  
  5472  	alloc6 := mock.Alloc()
  5473  	alloc6.JobID = alloc.JobID
  5474  	alloc6.Job = alloc.Job
  5475  	alloc6.TaskGroup = "db"
  5476  	alloc7 := alloc6.Copy()
  5477  	alloc7.ClientStatus = structs.AllocClientStatusComplete
  5478  
  5479  	alloc8 := mock.Alloc()
  5480  	alloc8.JobID = alloc.JobID
  5481  	alloc8.Job = alloc.Job
  5482  	alloc8.TaskGroup = "db"
  5483  	alloc9 := alloc8.Copy()
  5484  	alloc9.ClientStatus = structs.AllocClientStatusFailed
  5485  
  5486  	alloc10 := mock.Alloc()
  5487  	alloc10.JobID = alloc.JobID
  5488  	alloc10.Job = alloc.Job
  5489  	alloc10.TaskGroup = "db"
  5490  	alloc11 := alloc10.Copy()
  5491  	alloc11.ClientStatus = structs.AllocClientStatusLost
  5492  
  5493  	state.UpsertAllocs(structs.MsgTypeTestSetup, 130, []*structs.Allocation{alloc4, alloc6, alloc8, alloc10})
  5494  
  5495  	state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 150, []*structs.Allocation{alloc5, alloc7, alloc9, alloc11})
  5496  
  5497  	// DeleteJobSummary is a helper method and doesn't modify the indexes table
  5498  	state.DeleteJobSummary(130, alloc.Namespace, alloc.Job.ID)
  5499  
  5500  	state.ReconcileJobSummaries(120)
  5501  
  5502  	ws := memdb.NewWatchSet()
  5503  	summary, _ := state.JobSummaryByID(ws, alloc.Namespace, alloc.Job.ID)
  5504  	expectedSummary := structs.JobSummary{
  5505  		JobID:     alloc.Job.ID,
  5506  		Namespace: alloc.Namespace,
  5507  		Summary: map[string]structs.TaskGroupSummary{
  5508  			"web": {
  5509  				Running: 1,
  5510  			},
  5511  			"db": {
  5512  				Starting: 1,
  5513  				Running:  1,
  5514  				Failed:   1,
  5515  				Complete: 1,
  5516  				Lost:     1,
  5517  			},
  5518  		},
  5519  		CreateIndex: 100,
  5520  		ModifyIndex: 120,
  5521  	}
  5522  	if !reflect.DeepEqual(&expectedSummary, summary) {
  5523  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  5524  	}
  5525  }
  5526  
  5527  func TestStateStore_ReconcileParentJobSummary(t *testing.T) {
  5528  	t.Parallel()
  5529  	require := require.New(t)
  5530  
  5531  	state := testStateStore(t)
  5532  
  5533  	// Add a node
  5534  	node := mock.Node()
  5535  	state.UpsertNode(structs.MsgTypeTestSetup, 80, node)
  5536  
  5537  	// Make a parameterized job
  5538  	job1 := mock.BatchJob()
  5539  	job1.ID = "test"
  5540  	job1.ParameterizedJob = &structs.ParameterizedJobConfig{
  5541  		Payload: "random",
  5542  	}
  5543  	job1.TaskGroups[0].Count = 1
  5544  	state.UpsertJob(structs.MsgTypeTestSetup, 100, job1)
  5545  
  5546  	// Make a child job
  5547  	childJob := job1.Copy()
  5548  	childJob.ID = job1.ID + "dispatch-23423423"
  5549  	childJob.ParentID = job1.ID
  5550  	childJob.Dispatched = true
  5551  	childJob.Status = structs.JobStatusRunning
  5552  
  5553  	// Make some allocs for child job
  5554  	alloc := mock.Alloc()
  5555  	alloc.NodeID = node.ID
  5556  	alloc.Job = childJob
  5557  	alloc.JobID = childJob.ID
  5558  	alloc.ClientStatus = structs.AllocClientStatusRunning
  5559  
  5560  	alloc2 := mock.Alloc()
  5561  	alloc2.NodeID = node.ID
  5562  	alloc2.Job = childJob
  5563  	alloc2.JobID = childJob.ID
  5564  	alloc2.ClientStatus = structs.AllocClientStatusFailed
  5565  
  5566  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 110, childJob))
  5567  	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 111, []*structs.Allocation{alloc, alloc2}))
  5568  
  5569  	// Make the summary incorrect in the state store
  5570  	summary, err := state.JobSummaryByID(nil, job1.Namespace, job1.ID)
  5571  	require.Nil(err)
  5572  
  5573  	summary.Children = nil
  5574  	summary.Summary = make(map[string]structs.TaskGroupSummary)
  5575  	summary.Summary["web"] = structs.TaskGroupSummary{
  5576  		Queued: 1,
  5577  	}
  5578  
  5579  	// Delete the child job summary
  5580  	state.DeleteJobSummary(125, childJob.Namespace, childJob.ID)
  5581  
  5582  	state.ReconcileJobSummaries(120)
  5583  
  5584  	ws := memdb.NewWatchSet()
  5585  
  5586  	// Verify parent summary is corrected
  5587  	summary, _ = state.JobSummaryByID(ws, alloc.Namespace, job1.ID)
  5588  	expectedSummary := structs.JobSummary{
  5589  		JobID:     job1.ID,
  5590  		Namespace: job1.Namespace,
  5591  		Summary:   make(map[string]structs.TaskGroupSummary),
  5592  		Children: &structs.JobChildrenSummary{
  5593  			Running: 1,
  5594  		},
  5595  		CreateIndex: 100,
  5596  		ModifyIndex: 120,
  5597  	}
  5598  	require.Equal(&expectedSummary, summary)
  5599  
  5600  	// Verify child job summary is also correct
  5601  	childSummary, _ := state.JobSummaryByID(ws, childJob.Namespace, childJob.ID)
  5602  	expectedChildSummary := structs.JobSummary{
  5603  		JobID:     childJob.ID,
  5604  		Namespace: childJob.Namespace,
  5605  		Summary: map[string]structs.TaskGroupSummary{
  5606  			"web": {
  5607  				Running: 1,
  5608  				Failed:  1,
  5609  			},
  5610  		},
  5611  		CreateIndex: 110,
  5612  		ModifyIndex: 120,
  5613  	}
  5614  	require.Equal(&expectedChildSummary, childSummary)
  5615  }
  5616  
  5617  func TestStateStore_UpdateAlloc_JobNotPresent(t *testing.T) {
  5618  	t.Parallel()
  5619  
  5620  	state := testStateStore(t)
  5621  
  5622  	alloc := mock.Alloc()
  5623  	state.UpsertJob(structs.MsgTypeTestSetup, 100, alloc.Job)
  5624  	state.UpsertAllocs(structs.MsgTypeTestSetup, 200, []*structs.Allocation{alloc})
  5625  
  5626  	// Delete the job
  5627  	state.DeleteJob(300, alloc.Namespace, alloc.Job.ID)
  5628  
  5629  	// Update the alloc
  5630  	alloc1 := alloc.Copy()
  5631  	alloc1.ClientStatus = structs.AllocClientStatusRunning
  5632  
  5633  	// Updating allocation should not throw any error
  5634  	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 400, []*structs.Allocation{alloc1}); err != nil {
  5635  		t.Fatalf("expect err: %v", err)
  5636  	}
  5637  
  5638  	// Re-Register the job
  5639  	state.UpsertJob(structs.MsgTypeTestSetup, 500, alloc.Job)
  5640  
  5641  	// Update the alloc again
  5642  	alloc2 := alloc.Copy()
  5643  	alloc2.ClientStatus = structs.AllocClientStatusComplete
  5644  	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 400, []*structs.Allocation{alloc1}); err != nil {
  5645  		t.Fatalf("expect err: %v", err)
  5646  	}
  5647  
  5648  	// Job Summary of the newly registered job shouldn't account for the
  5649  	// allocation update for the older job
  5650  	expectedSummary := structs.JobSummary{
  5651  		JobID:     alloc1.JobID,
  5652  		Namespace: alloc1.Namespace,
  5653  		Summary: map[string]structs.TaskGroupSummary{
  5654  			"web": {},
  5655  		},
  5656  		Children:    new(structs.JobChildrenSummary),
  5657  		CreateIndex: 500,
  5658  		ModifyIndex: 500,
  5659  	}
  5660  
  5661  	ws := memdb.NewWatchSet()
  5662  	summary, _ := state.JobSummaryByID(ws, alloc.Namespace, alloc.Job.ID)
  5663  	if !reflect.DeepEqual(&expectedSummary, summary) {
  5664  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  5665  	}
  5666  }
  5667  
  5668  func TestStateStore_EvictAlloc_Alloc(t *testing.T) {
  5669  	t.Parallel()
  5670  
  5671  	state := testStateStore(t)
  5672  	alloc := mock.Alloc()
  5673  
  5674  	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
  5675  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc})
  5676  	if err != nil {
  5677  		t.Fatalf("err: %v", err)
  5678  	}
  5679  
  5680  	alloc2 := new(structs.Allocation)
  5681  	*alloc2 = *alloc
  5682  	alloc2.DesiredStatus = structs.AllocDesiredStatusEvict
  5683  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc2})
  5684  	if err != nil {
  5685  		t.Fatalf("err: %v", err)
  5686  	}
  5687  
  5688  	ws := memdb.NewWatchSet()
  5689  	out, err := state.AllocByID(ws, alloc.ID)
  5690  	if err != nil {
  5691  		t.Fatalf("err: %v", err)
  5692  	}
  5693  
  5694  	if out.DesiredStatus != structs.AllocDesiredStatusEvict {
  5695  		t.Fatalf("bad: %#v %#v", alloc, out)
  5696  	}
  5697  
  5698  	index, err := state.Index("allocs")
  5699  	if err != nil {
  5700  		t.Fatalf("err: %v", err)
  5701  	}
  5702  	if index != 1001 {
  5703  		t.Fatalf("bad: %d", index)
  5704  	}
  5705  }
  5706  
  5707  func TestStateStore_AllocsByNode(t *testing.T) {
  5708  	t.Parallel()
  5709  
  5710  	state := testStateStore(t)
  5711  	var allocs []*structs.Allocation
  5712  
  5713  	for i := 0; i < 10; i++ {
  5714  		alloc := mock.Alloc()
  5715  		alloc.NodeID = "foo"
  5716  		allocs = append(allocs, alloc)
  5717  	}
  5718  
  5719  	for idx, alloc := range allocs {
  5720  		state.UpsertJobSummary(uint64(900+idx), mock.JobSummary(alloc.JobID))
  5721  	}
  5722  
  5723  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
  5724  	if err != nil {
  5725  		t.Fatalf("err: %v", err)
  5726  	}
  5727  
  5728  	ws := memdb.NewWatchSet()
  5729  	out, err := state.AllocsByNode(ws, "foo")
  5730  	if err != nil {
  5731  		t.Fatalf("err: %v", err)
  5732  	}
  5733  
  5734  	sort.Sort(AllocIDSort(allocs))
  5735  	sort.Sort(AllocIDSort(out))
  5736  
  5737  	if !reflect.DeepEqual(allocs, out) {
  5738  		t.Fatalf("bad: %#v %#v", allocs, out)
  5739  	}
  5740  
  5741  	if watchFired(ws) {
  5742  		t.Fatalf("bad")
  5743  	}
  5744  }
  5745  
  5746  func TestStateStore_AllocsByNodeTerminal(t *testing.T) {
  5747  	t.Parallel()
  5748  
  5749  	state := testStateStore(t)
  5750  	var allocs, term, nonterm []*structs.Allocation
  5751  
  5752  	for i := 0; i < 10; i++ {
  5753  		alloc := mock.Alloc()
  5754  		alloc.NodeID = "foo"
  5755  		if i%2 == 0 {
  5756  			alloc.DesiredStatus = structs.AllocDesiredStatusStop
  5757  			term = append(term, alloc)
  5758  		} else {
  5759  			nonterm = append(nonterm, alloc)
  5760  		}
  5761  		allocs = append(allocs, alloc)
  5762  	}
  5763  
  5764  	for idx, alloc := range allocs {
  5765  		state.UpsertJobSummary(uint64(900+idx), mock.JobSummary(alloc.JobID))
  5766  	}
  5767  
  5768  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
  5769  	if err != nil {
  5770  		t.Fatalf("err: %v", err)
  5771  	}
  5772  
  5773  	// Verify the terminal allocs
  5774  	ws := memdb.NewWatchSet()
  5775  	out, err := state.AllocsByNodeTerminal(ws, "foo", true)
  5776  	if err != nil {
  5777  		t.Fatalf("err: %v", err)
  5778  	}
  5779  
  5780  	sort.Sort(AllocIDSort(term))
  5781  	sort.Sort(AllocIDSort(out))
  5782  
  5783  	if !reflect.DeepEqual(term, out) {
  5784  		t.Fatalf("bad: %#v %#v", term, out)
  5785  	}
  5786  
  5787  	// Verify the non-terminal allocs
  5788  	out, err = state.AllocsByNodeTerminal(ws, "foo", false)
  5789  	if err != nil {
  5790  		t.Fatalf("err: %v", err)
  5791  	}
  5792  
  5793  	sort.Sort(AllocIDSort(nonterm))
  5794  	sort.Sort(AllocIDSort(out))
  5795  
  5796  	if !reflect.DeepEqual(nonterm, out) {
  5797  		t.Fatalf("bad: %#v %#v", nonterm, out)
  5798  	}
  5799  
  5800  	if watchFired(ws) {
  5801  		t.Fatalf("bad")
  5802  	}
  5803  }
  5804  
  5805  func TestStateStore_AllocsByJob(t *testing.T) {
  5806  	t.Parallel()
  5807  
  5808  	state := testStateStore(t)
  5809  	var allocs []*structs.Allocation
  5810  
  5811  	for i := 0; i < 10; i++ {
  5812  		alloc := mock.Alloc()
  5813  		alloc.JobID = "foo"
  5814  		allocs = append(allocs, alloc)
  5815  	}
  5816  
  5817  	for i, alloc := range allocs {
  5818  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  5819  	}
  5820  
  5821  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
  5822  	if err != nil {
  5823  		t.Fatalf("err: %v", err)
  5824  	}
  5825  
  5826  	ws := memdb.NewWatchSet()
  5827  	out, err := state.AllocsByJob(ws, mock.Alloc().Namespace, "foo", false)
  5828  	if err != nil {
  5829  		t.Fatalf("err: %v", err)
  5830  	}
  5831  
  5832  	sort.Sort(AllocIDSort(allocs))
  5833  	sort.Sort(AllocIDSort(out))
  5834  
  5835  	if !reflect.DeepEqual(allocs, out) {
  5836  		t.Fatalf("bad: %#v %#v", allocs, out)
  5837  	}
  5838  
  5839  	if watchFired(ws) {
  5840  		t.Fatalf("bad")
  5841  	}
  5842  }
  5843  
  5844  func TestStateStore_AllocsForRegisteredJob(t *testing.T) {
  5845  	t.Parallel()
  5846  
  5847  	state := testStateStore(t)
  5848  	var allocs []*structs.Allocation
  5849  	var allocs1 []*structs.Allocation
  5850  
  5851  	job := mock.Job()
  5852  	job.ID = "foo"
  5853  	state.UpsertJob(structs.MsgTypeTestSetup, 100, job)
  5854  	for i := 0; i < 3; i++ {
  5855  		alloc := mock.Alloc()
  5856  		alloc.Job = job
  5857  		alloc.JobID = job.ID
  5858  		allocs = append(allocs, alloc)
  5859  	}
  5860  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 200, allocs); err != nil {
  5861  		t.Fatalf("err: %v", err)
  5862  	}
  5863  
  5864  	if err := state.DeleteJob(250, job.Namespace, job.ID); err != nil {
  5865  		t.Fatalf("err: %v", err)
  5866  	}
  5867  
  5868  	job1 := mock.Job()
  5869  	job1.ID = "foo"
  5870  	job1.CreateIndex = 50
  5871  	state.UpsertJob(structs.MsgTypeTestSetup, 300, job1)
  5872  	for i := 0; i < 4; i++ {
  5873  		alloc := mock.Alloc()
  5874  		alloc.Job = job1
  5875  		alloc.JobID = job1.ID
  5876  		allocs1 = append(allocs1, alloc)
  5877  	}
  5878  
  5879  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs1); err != nil {
  5880  		t.Fatalf("err: %v", err)
  5881  	}
  5882  
  5883  	ws := memdb.NewWatchSet()
  5884  	out, err := state.AllocsByJob(ws, job1.Namespace, job1.ID, true)
  5885  	if err != nil {
  5886  		t.Fatalf("err: %v", err)
  5887  	}
  5888  
  5889  	expected := len(allocs) + len(allocs1)
  5890  	if len(out) != expected {
  5891  		t.Fatalf("expected: %v, actual: %v", expected, len(out))
  5892  	}
  5893  
  5894  	out1, err := state.AllocsByJob(ws, job1.Namespace, job1.ID, false)
  5895  	if err != nil {
  5896  		t.Fatalf("bad: %v", err)
  5897  	}
  5898  
  5899  	expected = len(allocs1)
  5900  	if len(out1) != expected {
  5901  		t.Fatalf("expected: %v, actual: %v", expected, len(out1))
  5902  	}
  5903  
  5904  	if watchFired(ws) {
  5905  		t.Fatalf("bad")
  5906  	}
  5907  }
  5908  
  5909  func TestStateStore_AllocsByIDPrefix(t *testing.T) {
  5910  	t.Parallel()
  5911  
  5912  	state := testStateStore(t)
  5913  	var allocs []*structs.Allocation
  5914  
  5915  	ids := []string{
  5916  		"aaaaaaaa-7bfb-395d-eb95-0685af2176b2",
  5917  		"aaaaaaab-7bfb-395d-eb95-0685af2176b2",
  5918  		"aaaaaabb-7bfb-395d-eb95-0685af2176b2",
  5919  		"aaaaabbb-7bfb-395d-eb95-0685af2176b2",
  5920  		"aaaabbbb-7bfb-395d-eb95-0685af2176b2",
  5921  		"aaabbbbb-7bfb-395d-eb95-0685af2176b2",
  5922  		"aabbbbbb-7bfb-395d-eb95-0685af2176b2",
  5923  		"abbbbbbb-7bfb-395d-eb95-0685af2176b2",
  5924  		"bbbbbbbb-7bfb-395d-eb95-0685af2176b2",
  5925  	}
  5926  	for i := 0; i < 9; i++ {
  5927  		alloc := mock.Alloc()
  5928  		alloc.ID = ids[i]
  5929  		allocs = append(allocs, alloc)
  5930  	}
  5931  
  5932  	for i, alloc := range allocs {
  5933  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  5934  	}
  5935  
  5936  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
  5937  	if err != nil {
  5938  		t.Fatalf("err: %v", err)
  5939  	}
  5940  
  5941  	ws := memdb.NewWatchSet()
  5942  	iter, err := state.AllocsByIDPrefix(ws, structs.DefaultNamespace, "aaaa")
  5943  	if err != nil {
  5944  		t.Fatalf("err: %v", err)
  5945  	}
  5946  
  5947  	gatherAllocs := func(iter memdb.ResultIterator) []*structs.Allocation {
  5948  		var allocs []*structs.Allocation
  5949  		for {
  5950  			raw := iter.Next()
  5951  			if raw == nil {
  5952  				break
  5953  			}
  5954  			allocs = append(allocs, raw.(*structs.Allocation))
  5955  		}
  5956  		return allocs
  5957  	}
  5958  
  5959  	out := gatherAllocs(iter)
  5960  	if len(out) != 5 {
  5961  		t.Fatalf("bad: expected five allocations, got: %#v", out)
  5962  	}
  5963  
  5964  	sort.Sort(AllocIDSort(allocs))
  5965  
  5966  	for index, alloc := range out {
  5967  		if ids[index] != alloc.ID {
  5968  			t.Fatalf("bad: got unexpected id: %s", alloc.ID)
  5969  		}
  5970  	}
  5971  
  5972  	iter, err = state.AllocsByIDPrefix(ws, structs.DefaultNamespace, "b-a7bfb")
  5973  	if err != nil {
  5974  		t.Fatalf("err: %v", err)
  5975  	}
  5976  
  5977  	out = gatherAllocs(iter)
  5978  	if len(out) != 0 {
  5979  		t.Fatalf("bad: unexpected zero allocations, got: %#v", out)
  5980  	}
  5981  
  5982  	if watchFired(ws) {
  5983  		t.Fatalf("bad")
  5984  	}
  5985  }
  5986  
  5987  func TestStateStore_Allocs(t *testing.T) {
  5988  	t.Parallel()
  5989  
  5990  	state := testStateStore(t)
  5991  	var allocs []*structs.Allocation
  5992  
  5993  	for i := 0; i < 10; i++ {
  5994  		alloc := mock.Alloc()
  5995  		allocs = append(allocs, alloc)
  5996  	}
  5997  	for i, alloc := range allocs {
  5998  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  5999  	}
  6000  
  6001  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
  6002  	if err != nil {
  6003  		t.Fatalf("err: %v", err)
  6004  	}
  6005  
  6006  	ws := memdb.NewWatchSet()
  6007  	iter, err := state.Allocs(ws)
  6008  	if err != nil {
  6009  		t.Fatalf("err: %v", err)
  6010  	}
  6011  
  6012  	var out []*structs.Allocation
  6013  	for {
  6014  		raw := iter.Next()
  6015  		if raw == nil {
  6016  			break
  6017  		}
  6018  		out = append(out, raw.(*structs.Allocation))
  6019  	}
  6020  
  6021  	sort.Sort(AllocIDSort(allocs))
  6022  	sort.Sort(AllocIDSort(out))
  6023  
  6024  	if !reflect.DeepEqual(allocs, out) {
  6025  		t.Fatalf("bad: %#v %#v", allocs, out)
  6026  	}
  6027  
  6028  	if watchFired(ws) {
  6029  		t.Fatalf("bad")
  6030  	}
  6031  }
  6032  
  6033  func TestStateStore_Allocs_PrevAlloc(t *testing.T) {
  6034  	t.Parallel()
  6035  
  6036  	state := testStateStore(t)
  6037  	var allocs []*structs.Allocation
  6038  
  6039  	require := require.New(t)
  6040  	for i := 0; i < 5; i++ {
  6041  		alloc := mock.Alloc()
  6042  		allocs = append(allocs, alloc)
  6043  	}
  6044  	for i, alloc := range allocs {
  6045  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  6046  	}
  6047  	// Set some previous alloc ids
  6048  	allocs[1].PreviousAllocation = allocs[0].ID
  6049  	allocs[2].PreviousAllocation = allocs[1].ID
  6050  
  6051  	err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, allocs)
  6052  	require.Nil(err)
  6053  
  6054  	ws := memdb.NewWatchSet()
  6055  	iter, err := state.Allocs(ws)
  6056  	require.Nil(err)
  6057  
  6058  	var out []*structs.Allocation
  6059  	for {
  6060  		raw := iter.Next()
  6061  		if raw == nil {
  6062  			break
  6063  		}
  6064  		out = append(out, raw.(*structs.Allocation))
  6065  	}
  6066  
  6067  	// Set expected NextAllocation fields
  6068  	allocs[0].NextAllocation = allocs[1].ID
  6069  	allocs[1].NextAllocation = allocs[2].ID
  6070  
  6071  	sort.Sort(AllocIDSort(allocs))
  6072  	sort.Sort(AllocIDSort(out))
  6073  
  6074  	require.Equal(allocs, out)
  6075  	require.False(watchFired(ws))
  6076  
  6077  	// Insert another alloc, verify index of previous alloc also got updated
  6078  	alloc := mock.Alloc()
  6079  	alloc.PreviousAllocation = allocs[0].ID
  6080  	err = state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc})
  6081  	require.Nil(err)
  6082  	alloc0, err := state.AllocByID(nil, allocs[0].ID)
  6083  	require.Nil(err)
  6084  	require.Equal(alloc0.ModifyIndex, uint64(1001))
  6085  }
  6086  
  6087  func TestStateStore_RestoreAlloc(t *testing.T) {
  6088  	t.Parallel()
  6089  
  6090  	state := testStateStore(t)
  6091  	alloc := mock.Alloc()
  6092  
  6093  	restore, err := state.Restore()
  6094  	if err != nil {
  6095  		t.Fatalf("err: %v", err)
  6096  	}
  6097  
  6098  	err = restore.AllocRestore(alloc)
  6099  	if err != nil {
  6100  		t.Fatalf("err: %v", err)
  6101  	}
  6102  
  6103  	require.NoError(t, restore.Commit())
  6104  
  6105  	ws := memdb.NewWatchSet()
  6106  	out, err := state.AllocByID(ws, alloc.ID)
  6107  	if err != nil {
  6108  		t.Fatalf("err: %v", err)
  6109  	}
  6110  
  6111  	if !reflect.DeepEqual(out, alloc) {
  6112  		t.Fatalf("Bad: %#v %#v", out, alloc)
  6113  	}
  6114  
  6115  	if watchFired(ws) {
  6116  		t.Fatalf("bad")
  6117  	}
  6118  }
  6119  
  6120  func TestStateStore_SetJobStatus_ForceStatus(t *testing.T) {
  6121  	t.Parallel()
  6122  
  6123  	index := uint64(0)
  6124  	state := testStateStore(t)
  6125  	txn := state.db.WriteTxn(index)
  6126  
  6127  	// Create and insert a mock job.
  6128  	job := mock.Job()
  6129  	job.Status = ""
  6130  	job.ModifyIndex = index
  6131  	if err := txn.Insert("jobs", job); err != nil {
  6132  		t.Fatalf("job insert failed: %v", err)
  6133  	}
  6134  
  6135  	exp := "foobar"
  6136  	index = uint64(1000)
  6137  	if err := state.setJobStatus(index, txn, job, false, exp); err != nil {
  6138  		t.Fatalf("setJobStatus() failed: %v", err)
  6139  	}
  6140  
  6141  	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
  6142  	if err != nil {
  6143  		t.Fatalf("job lookup failed: %v", err)
  6144  	}
  6145  	updated := i.(*structs.Job)
  6146  
  6147  	if updated.Status != exp {
  6148  		t.Fatalf("setJobStatus() set %v; expected %v", updated.Status, exp)
  6149  	}
  6150  
  6151  	if updated.ModifyIndex != index {
  6152  		t.Fatalf("setJobStatus() set %d; expected %d", updated.ModifyIndex, index)
  6153  	}
  6154  }
  6155  
  6156  func TestStateStore_SetJobStatus_NoOp(t *testing.T) {
  6157  	t.Parallel()
  6158  
  6159  	index := uint64(0)
  6160  	state := testStateStore(t)
  6161  	txn := state.db.WriteTxn(index)
  6162  
  6163  	// Create and insert a mock job that should be pending.
  6164  	job := mock.Job()
  6165  	job.Status = structs.JobStatusPending
  6166  	job.ModifyIndex = 10
  6167  	if err := txn.Insert("jobs", job); err != nil {
  6168  		t.Fatalf("job insert failed: %v", err)
  6169  	}
  6170  
  6171  	index = uint64(1000)
  6172  	if err := state.setJobStatus(index, txn, job, false, ""); err != nil {
  6173  		t.Fatalf("setJobStatus() failed: %v", err)
  6174  	}
  6175  
  6176  	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
  6177  	if err != nil {
  6178  		t.Fatalf("job lookup failed: %v", err)
  6179  	}
  6180  	updated := i.(*structs.Job)
  6181  
  6182  	if updated.ModifyIndex == index {
  6183  		t.Fatalf("setJobStatus() should have been a no-op")
  6184  	}
  6185  }
  6186  
  6187  func TestStateStore_SetJobStatus(t *testing.T) {
  6188  	t.Parallel()
  6189  
  6190  	state := testStateStore(t)
  6191  	txn := state.db.WriteTxn(uint64(0))
  6192  
  6193  	// Create and insert a mock job that should be pending but has an incorrect
  6194  	// status.
  6195  	job := mock.Job()
  6196  	job.Status = "foobar"
  6197  	job.ModifyIndex = 10
  6198  	if err := txn.Insert("jobs", job); err != nil {
  6199  		t.Fatalf("job insert failed: %v", err)
  6200  	}
  6201  
  6202  	index := uint64(1000)
  6203  	if err := state.setJobStatus(index, txn, job, false, ""); err != nil {
  6204  		t.Fatalf("setJobStatus() failed: %v", err)
  6205  	}
  6206  
  6207  	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
  6208  	if err != nil {
  6209  		t.Fatalf("job lookup failed: %v", err)
  6210  	}
  6211  	updated := i.(*structs.Job)
  6212  
  6213  	if updated.Status != structs.JobStatusPending {
  6214  		t.Fatalf("setJobStatus() set %v; expected %v", updated.Status, structs.JobStatusPending)
  6215  	}
  6216  
  6217  	if updated.ModifyIndex != index {
  6218  		t.Fatalf("setJobStatus() set %d; expected %d", updated.ModifyIndex, index)
  6219  	}
  6220  }
  6221  
  6222  func TestStateStore_GetJobStatus_NoEvalsOrAllocs(t *testing.T) {
  6223  	t.Parallel()
  6224  
  6225  	job := mock.Job()
  6226  	state := testStateStore(t)
  6227  	txn := state.db.ReadTxn()
  6228  	status, err := state.getJobStatus(txn, job, false)
  6229  	if err != nil {
  6230  		t.Fatalf("getJobStatus() failed: %v", err)
  6231  	}
  6232  
  6233  	if status != structs.JobStatusPending {
  6234  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusPending)
  6235  	}
  6236  }
  6237  
  6238  func TestStateStore_GetJobStatus_NoEvalsOrAllocs_Periodic(t *testing.T) {
  6239  	t.Parallel()
  6240  
  6241  	job := mock.PeriodicJob()
  6242  	state := testStateStore(t)
  6243  	txn := state.db.ReadTxn()
  6244  	status, err := state.getJobStatus(txn, job, false)
  6245  	if err != nil {
  6246  		t.Fatalf("getJobStatus() failed: %v", err)
  6247  	}
  6248  
  6249  	if status != structs.JobStatusRunning {
  6250  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  6251  	}
  6252  }
  6253  
  6254  func TestStateStore_GetJobStatus_NoEvalsOrAllocs_EvalDelete(t *testing.T) {
  6255  	t.Parallel()
  6256  
  6257  	job := mock.Job()
  6258  	state := testStateStore(t)
  6259  	txn := state.db.ReadTxn()
  6260  	status, err := state.getJobStatus(txn, job, true)
  6261  	if err != nil {
  6262  		t.Fatalf("getJobStatus() failed: %v", err)
  6263  	}
  6264  
  6265  	if status != structs.JobStatusDead {
  6266  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  6267  	}
  6268  }
  6269  
  6270  func TestStateStore_GetJobStatus_DeadEvalsAndAllocs(t *testing.T) {
  6271  	t.Parallel()
  6272  
  6273  	state := testStateStore(t)
  6274  	job := mock.Job()
  6275  
  6276  	// Create a mock alloc that is dead.
  6277  	alloc := mock.Alloc()
  6278  	alloc.JobID = job.ID
  6279  	alloc.DesiredStatus = structs.AllocDesiredStatusStop
  6280  	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
  6281  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc}); err != nil {
  6282  		t.Fatalf("err: %v", err)
  6283  	}
  6284  
  6285  	// Create a mock eval that is complete
  6286  	eval := mock.Eval()
  6287  	eval.JobID = job.ID
  6288  	eval.Status = structs.EvalStatusComplete
  6289  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1001, []*structs.Evaluation{eval}); err != nil {
  6290  		t.Fatalf("err: %v", err)
  6291  	}
  6292  
  6293  	txn := state.db.ReadTxn()
  6294  	status, err := state.getJobStatus(txn, job, false)
  6295  	if err != nil {
  6296  		t.Fatalf("getJobStatus() failed: %v", err)
  6297  	}
  6298  
  6299  	if status != structs.JobStatusDead {
  6300  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  6301  	}
  6302  }
  6303  
  6304  func TestStateStore_GetJobStatus_RunningAlloc(t *testing.T) {
  6305  	t.Parallel()
  6306  
  6307  	state := testStateStore(t)
  6308  	job := mock.Job()
  6309  
  6310  	// Create a mock alloc that is running.
  6311  	alloc := mock.Alloc()
  6312  	alloc.JobID = job.ID
  6313  	alloc.DesiredStatus = structs.AllocDesiredStatusRun
  6314  	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
  6315  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1000, []*structs.Allocation{alloc}); err != nil {
  6316  		t.Fatalf("err: %v", err)
  6317  	}
  6318  
  6319  	txn := state.db.ReadTxn()
  6320  	status, err := state.getJobStatus(txn, job, true)
  6321  	if err != nil {
  6322  		t.Fatalf("getJobStatus() failed: %v", err)
  6323  	}
  6324  
  6325  	if status != structs.JobStatusRunning {
  6326  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  6327  	}
  6328  }
  6329  
  6330  func TestStateStore_GetJobStatus_PeriodicJob(t *testing.T) {
  6331  	t.Parallel()
  6332  
  6333  	state := testStateStore(t)
  6334  	job := mock.PeriodicJob()
  6335  
  6336  	txn := state.db.ReadTxn()
  6337  	status, err := state.getJobStatus(txn, job, false)
  6338  	if err != nil {
  6339  		t.Fatalf("getJobStatus() failed: %v", err)
  6340  	}
  6341  
  6342  	if status != structs.JobStatusRunning {
  6343  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  6344  	}
  6345  
  6346  	// Mark it as stopped
  6347  	job.Stop = true
  6348  	status, err = state.getJobStatus(txn, job, false)
  6349  	if err != nil {
  6350  		t.Fatalf("getJobStatus() failed: %v", err)
  6351  	}
  6352  
  6353  	if status != structs.JobStatusDead {
  6354  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  6355  	}
  6356  }
  6357  
  6358  func TestStateStore_GetJobStatus_ParameterizedJob(t *testing.T) {
  6359  	t.Parallel()
  6360  
  6361  	state := testStateStore(t)
  6362  	job := mock.Job()
  6363  	job.ParameterizedJob = &structs.ParameterizedJobConfig{}
  6364  
  6365  	txn := state.db.ReadTxn()
  6366  	status, err := state.getJobStatus(txn, job, false)
  6367  	if err != nil {
  6368  		t.Fatalf("getJobStatus() failed: %v", err)
  6369  	}
  6370  
  6371  	if status != structs.JobStatusRunning {
  6372  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  6373  	}
  6374  
  6375  	// Mark it as stopped
  6376  	job.Stop = true
  6377  	status, err = state.getJobStatus(txn, job, false)
  6378  	if err != nil {
  6379  		t.Fatalf("getJobStatus() failed: %v", err)
  6380  	}
  6381  
  6382  	if status != structs.JobStatusDead {
  6383  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  6384  	}
  6385  }
  6386  
  6387  func TestStateStore_SetJobStatus_PendingEval(t *testing.T) {
  6388  	t.Parallel()
  6389  
  6390  	state := testStateStore(t)
  6391  	job := mock.Job()
  6392  
  6393  	// Create a mock eval that is pending.
  6394  	eval := mock.Eval()
  6395  	eval.JobID = job.ID
  6396  	eval.Status = structs.EvalStatusPending
  6397  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval}); err != nil {
  6398  		t.Fatalf("err: %v", err)
  6399  	}
  6400  
  6401  	txn := state.db.ReadTxn()
  6402  	status, err := state.getJobStatus(txn, job, true)
  6403  	if err != nil {
  6404  		t.Fatalf("getJobStatus() failed: %v", err)
  6405  	}
  6406  
  6407  	if status != structs.JobStatusPending {
  6408  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusPending)
  6409  	}
  6410  }
  6411  
  6412  // TestStateStore_SetJobStatus_SystemJob asserts that system jobs are still
  6413  // considered running until explicitly stopped.
  6414  func TestStateStore_SetJobStatus_SystemJob(t *testing.T) {
  6415  	t.Parallel()
  6416  
  6417  	state := testStateStore(t)
  6418  	job := mock.SystemJob()
  6419  
  6420  	// Create a mock eval that is pending.
  6421  	eval := mock.Eval()
  6422  	eval.JobID = job.ID
  6423  	eval.Type = job.Type
  6424  	eval.Status = structs.EvalStatusComplete
  6425  	if err := state.UpsertEvals(structs.MsgTypeTestSetup, 1000, []*structs.Evaluation{eval}); err != nil {
  6426  		t.Fatalf("err: %v", err)
  6427  	}
  6428  
  6429  	txn := state.db.ReadTxn()
  6430  	status, err := state.getJobStatus(txn, job, true)
  6431  	if err != nil {
  6432  		t.Fatalf("getJobStatus() failed: %v", err)
  6433  	}
  6434  
  6435  	if expected := structs.JobStatusRunning; status != expected {
  6436  		t.Fatalf("getJobStatus() returned %v; expected %v", status, expected)
  6437  	}
  6438  
  6439  	// Stop the job
  6440  	job.Stop = true
  6441  	status, err = state.getJobStatus(txn, job, true)
  6442  	if err != nil {
  6443  		t.Fatalf("getJobStatus() failed: %v", err)
  6444  	}
  6445  
  6446  	if expected := structs.JobStatusDead; status != expected {
  6447  		t.Fatalf("getJobStatus() returned %v; expected %v", status, expected)
  6448  	}
  6449  }
  6450  
  6451  func TestStateJobSummary_UpdateJobCount(t *testing.T) {
  6452  	t.Parallel()
  6453  
  6454  	state := testStateStore(t)
  6455  	alloc := mock.Alloc()
  6456  	job := alloc.Job
  6457  	job.TaskGroups[0].Count = 3
  6458  
  6459  	// Create watchsets so we can test that upsert fires the watch
  6460  	ws := memdb.NewWatchSet()
  6461  	if _, err := state.JobSummaryByID(ws, job.Namespace, job.ID); err != nil {
  6462  		t.Fatalf("bad: %v", err)
  6463  	}
  6464  
  6465  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job); err != nil {
  6466  		t.Fatalf("err: %v", err)
  6467  	}
  6468  
  6469  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc}); err != nil {
  6470  		t.Fatalf("err: %v", err)
  6471  	}
  6472  
  6473  	if !watchFired(ws) {
  6474  		t.Fatalf("bad")
  6475  	}
  6476  
  6477  	ws = memdb.NewWatchSet()
  6478  	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
  6479  	expectedSummary := structs.JobSummary{
  6480  		JobID:     job.ID,
  6481  		Namespace: job.Namespace,
  6482  		Summary: map[string]structs.TaskGroupSummary{
  6483  			"web": {
  6484  				Starting: 1,
  6485  			},
  6486  		},
  6487  		Children:    new(structs.JobChildrenSummary),
  6488  		CreateIndex: 1000,
  6489  		ModifyIndex: 1001,
  6490  	}
  6491  	if !reflect.DeepEqual(summary, &expectedSummary) {
  6492  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  6493  	}
  6494  
  6495  	// Create watchsets so we can test that upsert fires the watch
  6496  	ws2 := memdb.NewWatchSet()
  6497  	if _, err := state.JobSummaryByID(ws2, job.Namespace, job.ID); err != nil {
  6498  		t.Fatalf("bad: %v", err)
  6499  	}
  6500  
  6501  	alloc2 := mock.Alloc()
  6502  	alloc2.Job = job
  6503  	alloc2.JobID = job.ID
  6504  
  6505  	alloc3 := mock.Alloc()
  6506  	alloc3.Job = job
  6507  	alloc3.JobID = job.ID
  6508  
  6509  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{alloc2, alloc3}); err != nil {
  6510  		t.Fatalf("err: %v", err)
  6511  	}
  6512  
  6513  	if !watchFired(ws2) {
  6514  		t.Fatalf("bad")
  6515  	}
  6516  
  6517  	outA, _ := state.AllocByID(ws, alloc3.ID)
  6518  
  6519  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  6520  	expectedSummary = structs.JobSummary{
  6521  		JobID:     job.ID,
  6522  		Namespace: job.Namespace,
  6523  		Summary: map[string]structs.TaskGroupSummary{
  6524  			"web": {
  6525  				Starting: 3,
  6526  			},
  6527  		},
  6528  		Children:    new(structs.JobChildrenSummary),
  6529  		CreateIndex: job.CreateIndex,
  6530  		ModifyIndex: outA.ModifyIndex,
  6531  	}
  6532  	if !reflect.DeepEqual(summary, &expectedSummary) {
  6533  		t.Fatalf("expected summary: %v, actual: %v", expectedSummary, summary)
  6534  	}
  6535  
  6536  	// Create watchsets so we can test that upsert fires the watch
  6537  	ws3 := memdb.NewWatchSet()
  6538  	if _, err := state.JobSummaryByID(ws3, job.Namespace, job.ID); err != nil {
  6539  		t.Fatalf("bad: %v", err)
  6540  	}
  6541  
  6542  	alloc4 := mock.Alloc()
  6543  	alloc4.ID = alloc2.ID
  6544  	alloc4.Job = alloc2.Job
  6545  	alloc4.JobID = alloc2.JobID
  6546  	alloc4.ClientStatus = structs.AllocClientStatusComplete
  6547  
  6548  	alloc5 := mock.Alloc()
  6549  	alloc5.ID = alloc3.ID
  6550  	alloc5.Job = alloc3.Job
  6551  	alloc5.JobID = alloc3.JobID
  6552  	alloc5.ClientStatus = structs.AllocClientStatusComplete
  6553  
  6554  	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1004, []*structs.Allocation{alloc4, alloc5}); err != nil {
  6555  		t.Fatalf("err: %v", err)
  6556  	}
  6557  
  6558  	if !watchFired(ws2) {
  6559  		t.Fatalf("bad")
  6560  	}
  6561  
  6562  	outA, _ = state.AllocByID(ws, alloc5.ID)
  6563  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  6564  	expectedSummary = structs.JobSummary{
  6565  		JobID:     job.ID,
  6566  		Namespace: job.Namespace,
  6567  		Summary: map[string]structs.TaskGroupSummary{
  6568  			"web": {
  6569  				Complete: 2,
  6570  				Starting: 1,
  6571  			},
  6572  		},
  6573  		Children:    new(structs.JobChildrenSummary),
  6574  		CreateIndex: job.CreateIndex,
  6575  		ModifyIndex: outA.ModifyIndex,
  6576  	}
  6577  	if !reflect.DeepEqual(summary, &expectedSummary) {
  6578  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  6579  	}
  6580  }
  6581  
  6582  func TestJobSummary_UpdateClientStatus(t *testing.T) {
  6583  	t.Parallel()
  6584  
  6585  	state := testStateStore(t)
  6586  	alloc := mock.Alloc()
  6587  	job := alloc.Job
  6588  	job.TaskGroups[0].Count = 3
  6589  
  6590  	alloc2 := mock.Alloc()
  6591  	alloc2.Job = job
  6592  	alloc2.JobID = job.ID
  6593  
  6594  	alloc3 := mock.Alloc()
  6595  	alloc3.Job = job
  6596  	alloc3.JobID = job.ID
  6597  
  6598  	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  6599  	if err != nil {
  6600  		t.Fatalf("err: %v", err)
  6601  	}
  6602  
  6603  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1001, []*structs.Allocation{alloc, alloc2, alloc3}); err != nil {
  6604  		t.Fatalf("err: %v", err)
  6605  	}
  6606  
  6607  	ws := memdb.NewWatchSet()
  6608  	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
  6609  	if summary.Summary["web"].Starting != 3 {
  6610  		t.Fatalf("bad job summary: %v", summary)
  6611  	}
  6612  
  6613  	alloc4 := mock.Alloc()
  6614  	alloc4.ID = alloc2.ID
  6615  	alloc4.Job = alloc2.Job
  6616  	alloc4.JobID = alloc2.JobID
  6617  	alloc4.ClientStatus = structs.AllocClientStatusComplete
  6618  
  6619  	alloc5 := mock.Alloc()
  6620  	alloc5.ID = alloc3.ID
  6621  	alloc5.Job = alloc3.Job
  6622  	alloc5.JobID = alloc3.JobID
  6623  	alloc5.ClientStatus = structs.AllocClientStatusFailed
  6624  
  6625  	alloc6 := mock.Alloc()
  6626  	alloc6.ID = alloc.ID
  6627  	alloc6.Job = alloc.Job
  6628  	alloc6.JobID = alloc.JobID
  6629  	alloc6.ClientStatus = structs.AllocClientStatusRunning
  6630  
  6631  	if err := state.UpdateAllocsFromClient(structs.MsgTypeTestSetup, 1002, []*structs.Allocation{alloc4, alloc5, alloc6}); err != nil {
  6632  		t.Fatalf("err: %v", err)
  6633  	}
  6634  
  6635  	if !watchFired(ws) {
  6636  		t.Fatalf("bad")
  6637  	}
  6638  
  6639  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  6640  	if summary.Summary["web"].Running != 1 || summary.Summary["web"].Failed != 1 || summary.Summary["web"].Complete != 1 {
  6641  		t.Fatalf("bad job summary: %v", summary)
  6642  	}
  6643  
  6644  	alloc7 := mock.Alloc()
  6645  	alloc7.Job = alloc.Job
  6646  	alloc7.JobID = alloc.JobID
  6647  
  6648  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 1003, []*structs.Allocation{alloc7}); err != nil {
  6649  		t.Fatalf("err: %v", err)
  6650  	}
  6651  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  6652  	if summary.Summary["web"].Starting != 1 || summary.Summary["web"].Running != 1 || summary.Summary["web"].Failed != 1 || summary.Summary["web"].Complete != 1 {
  6653  		t.Fatalf("bad job summary: %v", summary)
  6654  	}
  6655  }
  6656  
  6657  // Test that nonexistent deployment can't be updated
  6658  func TestStateStore_UpsertDeploymentStatusUpdate_Nonexistent(t *testing.T) {
  6659  	t.Parallel()
  6660  
  6661  	state := testStateStore(t)
  6662  
  6663  	// Update the nonexistent deployment
  6664  	req := &structs.DeploymentStatusUpdateRequest{
  6665  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  6666  			DeploymentID: uuid.Generate(),
  6667  			Status:       structs.DeploymentStatusRunning,
  6668  		},
  6669  	}
  6670  	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 2, req)
  6671  	if err == nil || !strings.Contains(err.Error(), "does not exist") {
  6672  		t.Fatalf("expected error updating the status because the deployment doesn't exist")
  6673  	}
  6674  }
  6675  
  6676  // Test that terminal deployment can't be updated
  6677  func TestStateStore_UpsertDeploymentStatusUpdate_Terminal(t *testing.T) {
  6678  	t.Parallel()
  6679  
  6680  	state := testStateStore(t)
  6681  
  6682  	// Insert a terminal deployment
  6683  	d := mock.Deployment()
  6684  	d.Status = structs.DeploymentStatusFailed
  6685  
  6686  	if err := state.UpsertDeployment(1, d); err != nil {
  6687  		t.Fatalf("bad: %v", err)
  6688  	}
  6689  
  6690  	// Update the deployment
  6691  	req := &structs.DeploymentStatusUpdateRequest{
  6692  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  6693  			DeploymentID: d.ID,
  6694  			Status:       structs.DeploymentStatusRunning,
  6695  		},
  6696  	}
  6697  	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 2, req)
  6698  	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
  6699  		t.Fatalf("expected error updating the status because the deployment is terminal")
  6700  	}
  6701  }
  6702  
  6703  // Test that a non terminal deployment is updated and that a job and eval are
  6704  // created.
  6705  func TestStateStore_UpsertDeploymentStatusUpdate_NonTerminal(t *testing.T) {
  6706  	t.Parallel()
  6707  
  6708  	state := testStateStore(t)
  6709  
  6710  	// Insert a deployment
  6711  	d := mock.Deployment()
  6712  	if err := state.UpsertDeployment(1, d); err != nil {
  6713  		t.Fatalf("bad: %v", err)
  6714  	}
  6715  
  6716  	// Create an eval and a job
  6717  	e := mock.Eval()
  6718  	j := mock.Job()
  6719  
  6720  	// Update the deployment
  6721  	status, desc := structs.DeploymentStatusFailed, "foo"
  6722  	req := &structs.DeploymentStatusUpdateRequest{
  6723  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  6724  			DeploymentID:      d.ID,
  6725  			Status:            status,
  6726  			StatusDescription: desc,
  6727  		},
  6728  		Job:  j,
  6729  		Eval: e,
  6730  	}
  6731  	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 2, req)
  6732  	if err != nil {
  6733  		t.Fatalf("bad: %v", err)
  6734  	}
  6735  
  6736  	// Check that the status was updated properly
  6737  	ws := memdb.NewWatchSet()
  6738  	dout, err := state.DeploymentByID(ws, d.ID)
  6739  	if err != nil {
  6740  		t.Fatalf("bad: %v", err)
  6741  	}
  6742  	if dout.Status != status || dout.StatusDescription != desc {
  6743  		t.Fatalf("bad: %#v", dout)
  6744  	}
  6745  
  6746  	// Check that the evaluation was created
  6747  	eout, _ := state.EvalByID(ws, e.ID)
  6748  	if err != nil {
  6749  		t.Fatalf("bad: %v", err)
  6750  	}
  6751  	if eout == nil {
  6752  		t.Fatalf("bad: %#v", eout)
  6753  	}
  6754  
  6755  	// Check that the job was created
  6756  	jout, _ := state.JobByID(ws, j.Namespace, j.ID)
  6757  	if err != nil {
  6758  		t.Fatalf("bad: %v", err)
  6759  	}
  6760  	if jout == nil {
  6761  		t.Fatalf("bad: %#v", jout)
  6762  	}
  6763  }
  6764  
  6765  // Test that when a deployment is updated to successful the job is updated to
  6766  // stable
  6767  func TestStateStore_UpsertDeploymentStatusUpdate_Successful(t *testing.T) {
  6768  	t.Parallel()
  6769  
  6770  	state := testStateStore(t)
  6771  
  6772  	// Insert a job
  6773  	job := mock.Job()
  6774  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1, job); err != nil {
  6775  		t.Fatalf("bad: %v", err)
  6776  	}
  6777  
  6778  	// Insert a deployment
  6779  	d := structs.NewDeployment(job)
  6780  	if err := state.UpsertDeployment(2, d); err != nil {
  6781  		t.Fatalf("bad: %v", err)
  6782  	}
  6783  
  6784  	// Update the deployment
  6785  	req := &structs.DeploymentStatusUpdateRequest{
  6786  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  6787  			DeploymentID:      d.ID,
  6788  			Status:            structs.DeploymentStatusSuccessful,
  6789  			StatusDescription: structs.DeploymentStatusDescriptionSuccessful,
  6790  		},
  6791  	}
  6792  	err := state.UpdateDeploymentStatus(structs.MsgTypeTestSetup, 3, req)
  6793  	if err != nil {
  6794  		t.Fatalf("bad: %v", err)
  6795  	}
  6796  
  6797  	// Check that the status was updated properly
  6798  	ws := memdb.NewWatchSet()
  6799  	dout, err := state.DeploymentByID(ws, d.ID)
  6800  	if err != nil {
  6801  		t.Fatalf("bad: %v", err)
  6802  	}
  6803  	if dout.Status != structs.DeploymentStatusSuccessful ||
  6804  		dout.StatusDescription != structs.DeploymentStatusDescriptionSuccessful {
  6805  		t.Fatalf("bad: %#v", dout)
  6806  	}
  6807  
  6808  	// Check that the job was created
  6809  	jout, _ := state.JobByID(ws, job.Namespace, job.ID)
  6810  	if err != nil {
  6811  		t.Fatalf("bad: %v", err)
  6812  	}
  6813  	if jout == nil {
  6814  		t.Fatalf("bad: %#v", jout)
  6815  	}
  6816  	if !jout.Stable {
  6817  		t.Fatalf("job not marked stable %#v", jout)
  6818  	}
  6819  	if jout.Version != d.JobVersion {
  6820  		t.Fatalf("job version changed; got %d; want %d", jout.Version, d.JobVersion)
  6821  	}
  6822  }
  6823  
  6824  func TestStateStore_UpdateJobStability(t *testing.T) {
  6825  	t.Parallel()
  6826  
  6827  	state := testStateStore(t)
  6828  
  6829  	// Insert a job twice to get two versions
  6830  	job := mock.Job()
  6831  	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 1, job))
  6832  
  6833  	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 2, job.Copy()))
  6834  
  6835  	// Update the stability to true
  6836  	err := state.UpdateJobStability(3, job.Namespace, job.ID, 0, true)
  6837  	require.NoError(t, err)
  6838  
  6839  	// Check that the job was updated properly
  6840  	ws := memdb.NewWatchSet()
  6841  	jout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
  6842  	require.NoError(t, err)
  6843  	require.NotNil(t, jout)
  6844  	require.True(t, jout.Stable, "job not marked as stable")
  6845  
  6846  	// Update the stability to false
  6847  	err = state.UpdateJobStability(3, job.Namespace, job.ID, 0, false)
  6848  	if err != nil {
  6849  		t.Fatalf("bad: %v", err)
  6850  	}
  6851  
  6852  	// Check that the job was updated properly
  6853  	jout, err = state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
  6854  	require.NoError(t, err)
  6855  	require.NotNil(t, jout)
  6856  	require.False(t, jout.Stable)
  6857  }
  6858  
  6859  // Test that nonexistent deployment can't be promoted
  6860  func TestStateStore_UpsertDeploymentPromotion_Nonexistent(t *testing.T) {
  6861  	t.Parallel()
  6862  
  6863  	state := testStateStore(t)
  6864  
  6865  	// Promote the nonexistent deployment
  6866  	req := &structs.ApplyDeploymentPromoteRequest{
  6867  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  6868  			DeploymentID: uuid.Generate(),
  6869  			All:          true,
  6870  		},
  6871  	}
  6872  	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 2, req)
  6873  	if err == nil || !strings.Contains(err.Error(), "does not exist") {
  6874  		t.Fatalf("expected error promoting because the deployment doesn't exist")
  6875  	}
  6876  }
  6877  
  6878  // Test that terminal deployment can't be updated
  6879  func TestStateStore_UpsertDeploymentPromotion_Terminal(t *testing.T) {
  6880  	t.Parallel()
  6881  
  6882  	state := testStateStore(t)
  6883  
  6884  	// Insert a terminal deployment
  6885  	d := mock.Deployment()
  6886  	d.Status = structs.DeploymentStatusFailed
  6887  
  6888  	if err := state.UpsertDeployment(1, d); err != nil {
  6889  		t.Fatalf("bad: %v", err)
  6890  	}
  6891  
  6892  	// Promote the deployment
  6893  	req := &structs.ApplyDeploymentPromoteRequest{
  6894  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  6895  			DeploymentID: d.ID,
  6896  			All:          true,
  6897  		},
  6898  	}
  6899  	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 2, req)
  6900  	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
  6901  		t.Fatalf("expected error updating the status because the deployment is terminal: %v", err)
  6902  	}
  6903  }
  6904  
  6905  // Test promoting unhealthy canaries in a deployment.
  6906  func TestStateStore_UpsertDeploymentPromotion_Unhealthy(t *testing.T) {
  6907  	t.Parallel()
  6908  
  6909  	state := testStateStore(t)
  6910  	require := require.New(t)
  6911  
  6912  	// Create a job
  6913  	j := mock.Job()
  6914  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1, j))
  6915  
  6916  	// Create a deployment
  6917  	d := mock.Deployment()
  6918  	d.JobID = j.ID
  6919  	d.TaskGroups["web"].DesiredCanaries = 2
  6920  	require.Nil(state.UpsertDeployment(2, d))
  6921  
  6922  	// Create a set of allocations
  6923  	c1 := mock.Alloc()
  6924  	c1.JobID = j.ID
  6925  	c1.DeploymentID = d.ID
  6926  	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
  6927  	c2 := mock.Alloc()
  6928  	c2.JobID = j.ID
  6929  	c2.DeploymentID = d.ID
  6930  	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
  6931  
  6932  	// Create a healthy but terminal alloc
  6933  	c3 := mock.Alloc()
  6934  	c3.JobID = j.ID
  6935  	c3.DeploymentID = d.ID
  6936  	c3.DesiredStatus = structs.AllocDesiredStatusStop
  6937  	c3.DeploymentStatus = &structs.AllocDeploymentStatus{Healthy: helper.BoolToPtr(true)}
  6938  	d.TaskGroups[c3.TaskGroup].PlacedCanaries = append(d.TaskGroups[c3.TaskGroup].PlacedCanaries, c3.ID)
  6939  
  6940  	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{c1, c2, c3}))
  6941  
  6942  	// Promote the canaries
  6943  	req := &structs.ApplyDeploymentPromoteRequest{
  6944  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  6945  			DeploymentID: d.ID,
  6946  			All:          true,
  6947  		},
  6948  	}
  6949  	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req)
  6950  	require.NotNil(err)
  6951  	require.Contains(err.Error(), `Task group "web" has 0/2 healthy allocations`)
  6952  }
  6953  
  6954  // Test promoting a deployment with no canaries
  6955  func TestStateStore_UpsertDeploymentPromotion_NoCanaries(t *testing.T) {
  6956  	t.Parallel()
  6957  
  6958  	state := testStateStore(t)
  6959  	require := require.New(t)
  6960  
  6961  	// Create a job
  6962  	j := mock.Job()
  6963  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1, j))
  6964  
  6965  	// Create a deployment
  6966  	d := mock.Deployment()
  6967  	d.TaskGroups["web"].DesiredCanaries = 2
  6968  	d.JobID = j.ID
  6969  	require.Nil(state.UpsertDeployment(2, d))
  6970  
  6971  	// Promote the canaries
  6972  	req := &structs.ApplyDeploymentPromoteRequest{
  6973  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  6974  			DeploymentID: d.ID,
  6975  			All:          true,
  6976  		},
  6977  	}
  6978  	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req)
  6979  	require.NotNil(err)
  6980  	require.Contains(err.Error(), `Task group "web" has 0/2 healthy allocations`)
  6981  }
  6982  
  6983  // Test promoting all canaries in a deployment.
  6984  func TestStateStore_UpsertDeploymentPromotion_All(t *testing.T) {
  6985  	t.Parallel()
  6986  
  6987  	state := testStateStore(t)
  6988  
  6989  	// Create a job with two task groups
  6990  	j := mock.Job()
  6991  	tg1 := j.TaskGroups[0]
  6992  	tg2 := tg1.Copy()
  6993  	tg2.Name = "foo"
  6994  	j.TaskGroups = append(j.TaskGroups, tg2)
  6995  	if err := state.UpsertJob(structs.MsgTypeTestSetup, 1, j); err != nil {
  6996  		t.Fatalf("bad: %v", err)
  6997  	}
  6998  
  6999  	// Create a deployment
  7000  	d := mock.Deployment()
  7001  	d.StatusDescription = structs.DeploymentStatusDescriptionRunningNeedsPromotion
  7002  	d.JobID = j.ID
  7003  	d.TaskGroups = map[string]*structs.DeploymentState{
  7004  		"web": {
  7005  			DesiredTotal:    10,
  7006  			DesiredCanaries: 1,
  7007  		},
  7008  		"foo": {
  7009  			DesiredTotal:    10,
  7010  			DesiredCanaries: 1,
  7011  		},
  7012  	}
  7013  	if err := state.UpsertDeployment(2, d); err != nil {
  7014  		t.Fatalf("bad: %v", err)
  7015  	}
  7016  
  7017  	// Create a set of allocations
  7018  	c1 := mock.Alloc()
  7019  	c1.JobID = j.ID
  7020  	c1.DeploymentID = d.ID
  7021  	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
  7022  	c1.DeploymentStatus = &structs.AllocDeploymentStatus{
  7023  		Healthy: helper.BoolToPtr(true),
  7024  	}
  7025  	c2 := mock.Alloc()
  7026  	c2.JobID = j.ID
  7027  	c2.DeploymentID = d.ID
  7028  	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
  7029  	c2.TaskGroup = tg2.Name
  7030  	c2.DeploymentStatus = &structs.AllocDeploymentStatus{
  7031  		Healthy: helper.BoolToPtr(true),
  7032  	}
  7033  
  7034  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{c1, c2}); err != nil {
  7035  		t.Fatalf("err: %v", err)
  7036  	}
  7037  
  7038  	// Create an eval
  7039  	e := mock.Eval()
  7040  
  7041  	// Promote the canaries
  7042  	req := &structs.ApplyDeploymentPromoteRequest{
  7043  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  7044  			DeploymentID: d.ID,
  7045  			All:          true,
  7046  		},
  7047  		Eval: e,
  7048  	}
  7049  	err := state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req)
  7050  	if err != nil {
  7051  		t.Fatalf("bad: %v", err)
  7052  	}
  7053  
  7054  	// Check that the status per task group was updated properly
  7055  	ws := memdb.NewWatchSet()
  7056  	dout, err := state.DeploymentByID(ws, d.ID)
  7057  	if err != nil {
  7058  		t.Fatalf("bad: %v", err)
  7059  	}
  7060  	if dout.StatusDescription != structs.DeploymentStatusDescriptionRunning {
  7061  		t.Fatalf("status description not updated: got %v; want %v", dout.StatusDescription, structs.DeploymentStatusDescriptionRunning)
  7062  	}
  7063  	if len(dout.TaskGroups) != 2 {
  7064  		t.Fatalf("bad: %#v", dout.TaskGroups)
  7065  	}
  7066  	for tg, state := range dout.TaskGroups {
  7067  		if !state.Promoted {
  7068  			t.Fatalf("bad: group %q not promoted %#v", tg, state)
  7069  		}
  7070  	}
  7071  
  7072  	// Check that the evaluation was created
  7073  	eout, _ := state.EvalByID(ws, e.ID)
  7074  	if err != nil {
  7075  		t.Fatalf("bad: %v", err)
  7076  	}
  7077  	if eout == nil {
  7078  		t.Fatalf("bad: %#v", eout)
  7079  	}
  7080  }
  7081  
  7082  // Test promoting a subset of canaries in a deployment.
  7083  func TestStateStore_UpsertDeploymentPromotion_Subset(t *testing.T) {
  7084  	t.Parallel()
  7085  	require := require.New(t)
  7086  
  7087  	state := testStateStore(t)
  7088  
  7089  	// Create a job with two task groups
  7090  	j := mock.Job()
  7091  	tg1 := j.TaskGroups[0]
  7092  	tg2 := tg1.Copy()
  7093  	tg2.Name = "foo"
  7094  	j.TaskGroups = append(j.TaskGroups, tg2)
  7095  	require.Nil(state.UpsertJob(structs.MsgTypeTestSetup, 1, j))
  7096  
  7097  	// Create a deployment
  7098  	d := mock.Deployment()
  7099  	d.JobID = j.ID
  7100  	d.TaskGroups = map[string]*structs.DeploymentState{
  7101  		"web": {
  7102  			DesiredTotal:    10,
  7103  			DesiredCanaries: 1,
  7104  		},
  7105  		"foo": {
  7106  			DesiredTotal:    10,
  7107  			DesiredCanaries: 1,
  7108  		},
  7109  	}
  7110  	require.Nil(state.UpsertDeployment(2, d))
  7111  
  7112  	// Create a set of allocations for both groups, including an unhealthy one
  7113  	c1 := mock.Alloc()
  7114  	c1.JobID = j.ID
  7115  	c1.DeploymentID = d.ID
  7116  	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
  7117  	c1.DeploymentStatus = &structs.AllocDeploymentStatus{
  7118  		Healthy: helper.BoolToPtr(true),
  7119  		Canary:  true,
  7120  	}
  7121  
  7122  	// Should still be a canary
  7123  	c2 := mock.Alloc()
  7124  	c2.JobID = j.ID
  7125  	c2.DeploymentID = d.ID
  7126  	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
  7127  	c2.TaskGroup = tg2.Name
  7128  	c2.DeploymentStatus = &structs.AllocDeploymentStatus{
  7129  		Healthy: helper.BoolToPtr(true),
  7130  		Canary:  true,
  7131  	}
  7132  
  7133  	c3 := mock.Alloc()
  7134  	c3.JobID = j.ID
  7135  	c3.DeploymentID = d.ID
  7136  	d.TaskGroups[c3.TaskGroup].PlacedCanaries = append(d.TaskGroups[c3.TaskGroup].PlacedCanaries, c3.ID)
  7137  	c3.DeploymentStatus = &structs.AllocDeploymentStatus{
  7138  		Healthy: helper.BoolToPtr(false),
  7139  		Canary:  true,
  7140  	}
  7141  
  7142  	require.Nil(state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{c1, c2, c3}))
  7143  
  7144  	// Create an eval
  7145  	e := mock.Eval()
  7146  
  7147  	// Promote the canaries
  7148  	req := &structs.ApplyDeploymentPromoteRequest{
  7149  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  7150  			DeploymentID: d.ID,
  7151  			Groups:       []string{"web"},
  7152  		},
  7153  		Eval: e,
  7154  	}
  7155  	require.Nil(state.UpdateDeploymentPromotion(structs.MsgTypeTestSetup, 4, req))
  7156  
  7157  	// Check that the status per task group was updated properly
  7158  	ws := memdb.NewWatchSet()
  7159  	dout, err := state.DeploymentByID(ws, d.ID)
  7160  	require.Nil(err)
  7161  	require.Len(dout.TaskGroups, 2)
  7162  	require.Contains(dout.TaskGroups, "web")
  7163  	require.True(dout.TaskGroups["web"].Promoted)
  7164  
  7165  	// Check that the evaluation was created
  7166  	eout, err := state.EvalByID(ws, e.ID)
  7167  	require.Nil(err)
  7168  	require.NotNil(eout)
  7169  
  7170  	// Check the canary field was set properly
  7171  	aout1, err1 := state.AllocByID(ws, c1.ID)
  7172  	aout2, err2 := state.AllocByID(ws, c2.ID)
  7173  	aout3, err3 := state.AllocByID(ws, c3.ID)
  7174  	require.Nil(err1)
  7175  	require.Nil(err2)
  7176  	require.Nil(err3)
  7177  	require.NotNil(aout1)
  7178  	require.NotNil(aout2)
  7179  	require.NotNil(aout3)
  7180  	require.False(aout1.DeploymentStatus.Canary)
  7181  	require.True(aout2.DeploymentStatus.Canary)
  7182  	require.True(aout3.DeploymentStatus.Canary)
  7183  }
  7184  
  7185  // Test that allocation health can't be set against a nonexistent deployment
  7186  func TestStateStore_UpsertDeploymentAllocHealth_Nonexistent(t *testing.T) {
  7187  	t.Parallel()
  7188  
  7189  	state := testStateStore(t)
  7190  
  7191  	// Set health against the nonexistent deployment
  7192  	req := &structs.ApplyDeploymentAllocHealthRequest{
  7193  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  7194  			DeploymentID:         uuid.Generate(),
  7195  			HealthyAllocationIDs: []string{uuid.Generate()},
  7196  		},
  7197  	}
  7198  	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 2, req)
  7199  	if err == nil || !strings.Contains(err.Error(), "does not exist") {
  7200  		t.Fatalf("expected error because the deployment doesn't exist: %v", err)
  7201  	}
  7202  }
  7203  
  7204  // Test that allocation health can't be set against a terminal deployment
  7205  func TestStateStore_UpsertDeploymentAllocHealth_Terminal(t *testing.T) {
  7206  	t.Parallel()
  7207  
  7208  	state := testStateStore(t)
  7209  
  7210  	// Insert a terminal deployment
  7211  	d := mock.Deployment()
  7212  	d.Status = structs.DeploymentStatusFailed
  7213  
  7214  	if err := state.UpsertDeployment(1, d); err != nil {
  7215  		t.Fatalf("bad: %v", err)
  7216  	}
  7217  
  7218  	// Set health against the terminal deployment
  7219  	req := &structs.ApplyDeploymentAllocHealthRequest{
  7220  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  7221  			DeploymentID:         d.ID,
  7222  			HealthyAllocationIDs: []string{uuid.Generate()},
  7223  		},
  7224  	}
  7225  	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 2, req)
  7226  	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
  7227  		t.Fatalf("expected error because the deployment is terminal: %v", err)
  7228  	}
  7229  }
  7230  
  7231  // Test that allocation health can't be set against a nonexistent alloc
  7232  func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_Nonexistent(t *testing.T) {
  7233  	t.Parallel()
  7234  
  7235  	state := testStateStore(t)
  7236  
  7237  	// Insert a deployment
  7238  	d := mock.Deployment()
  7239  	if err := state.UpsertDeployment(1, d); err != nil {
  7240  		t.Fatalf("bad: %v", err)
  7241  	}
  7242  
  7243  	// Set health against the terminal deployment
  7244  	req := &structs.ApplyDeploymentAllocHealthRequest{
  7245  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  7246  			DeploymentID:         d.ID,
  7247  			HealthyAllocationIDs: []string{uuid.Generate()},
  7248  		},
  7249  	}
  7250  	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 2, req)
  7251  	if err == nil || !strings.Contains(err.Error(), "unknown alloc") {
  7252  		t.Fatalf("expected error because the alloc doesn't exist: %v", err)
  7253  	}
  7254  }
  7255  
  7256  // Test that a deployments PlacedCanaries is properly updated
  7257  func TestStateStore_UpsertDeploymentAlloc_Canaries(t *testing.T) {
  7258  	t.Parallel()
  7259  
  7260  	state := testStateStore(t)
  7261  
  7262  	// Create a deployment
  7263  	d1 := mock.Deployment()
  7264  	require.NoError(t, state.UpsertDeployment(2, d1))
  7265  
  7266  	// Create a Job
  7267  	job := mock.Job()
  7268  	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 3, job))
  7269  
  7270  	// Create alloc with canary status
  7271  	a := mock.Alloc()
  7272  	a.JobID = job.ID
  7273  	a.DeploymentID = d1.ID
  7274  	a.DeploymentStatus = &structs.AllocDeploymentStatus{
  7275  		Healthy: helper.BoolToPtr(false),
  7276  		Canary:  true,
  7277  	}
  7278  	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 4, []*structs.Allocation{a}))
  7279  
  7280  	// Pull the deployment from state
  7281  	ws := memdb.NewWatchSet()
  7282  	deploy, err := state.DeploymentByID(ws, d1.ID)
  7283  	require.NoError(t, err)
  7284  
  7285  	// Ensure that PlacedCanaries is accurate
  7286  	require.Equal(t, 1, len(deploy.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))
  7287  
  7288  	// Create alloc without canary status
  7289  	b := mock.Alloc()
  7290  	b.JobID = job.ID
  7291  	b.DeploymentID = d1.ID
  7292  	b.DeploymentStatus = &structs.AllocDeploymentStatus{
  7293  		Healthy: helper.BoolToPtr(false),
  7294  		Canary:  false,
  7295  	}
  7296  	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 4, []*structs.Allocation{b}))
  7297  
  7298  	// Pull the deployment from state
  7299  	ws = memdb.NewWatchSet()
  7300  	deploy, err = state.DeploymentByID(ws, d1.ID)
  7301  	require.NoError(t, err)
  7302  
  7303  	// Ensure that PlacedCanaries is accurate
  7304  	require.Equal(t, 1, len(deploy.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))
  7305  
  7306  	// Create a second deployment
  7307  	d2 := mock.Deployment()
  7308  	require.NoError(t, state.UpsertDeployment(5, d2))
  7309  
  7310  	c := mock.Alloc()
  7311  	c.JobID = job.ID
  7312  	c.DeploymentID = d2.ID
  7313  	c.DeploymentStatus = &structs.AllocDeploymentStatus{
  7314  		Healthy: helper.BoolToPtr(false),
  7315  		Canary:  true,
  7316  	}
  7317  	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 6, []*structs.Allocation{c}))
  7318  
  7319  	ws = memdb.NewWatchSet()
  7320  	deploy2, err := state.DeploymentByID(ws, d2.ID)
  7321  	require.NoError(t, err)
  7322  
  7323  	// Ensure that PlacedCanaries is accurate
  7324  	require.Equal(t, 1, len(deploy2.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))
  7325  }
  7326  
  7327  func TestStateStore_UpsertDeploymentAlloc_NoCanaries(t *testing.T) {
  7328  	t.Parallel()
  7329  
  7330  	state := testStateStore(t)
  7331  
  7332  	// Create a deployment
  7333  	d1 := mock.Deployment()
  7334  	require.NoError(t, state.UpsertDeployment(2, d1))
  7335  
  7336  	// Create a Job
  7337  	job := mock.Job()
  7338  	require.NoError(t, state.UpsertJob(structs.MsgTypeTestSetup, 3, job))
  7339  
  7340  	// Create alloc with canary status
  7341  	a := mock.Alloc()
  7342  	a.JobID = job.ID
  7343  	a.DeploymentID = d1.ID
  7344  	a.DeploymentStatus = &structs.AllocDeploymentStatus{
  7345  		Healthy: helper.BoolToPtr(true),
  7346  		Canary:  false,
  7347  	}
  7348  	require.NoError(t, state.UpsertAllocs(structs.MsgTypeTestSetup, 4, []*structs.Allocation{a}))
  7349  
  7350  	// Pull the deployment from state
  7351  	ws := memdb.NewWatchSet()
  7352  	deploy, err := state.DeploymentByID(ws, d1.ID)
  7353  	require.NoError(t, err)
  7354  
  7355  	// Ensure that PlacedCanaries is accurate
  7356  	require.Equal(t, 0, len(deploy.TaskGroups[job.TaskGroups[0].Name].PlacedCanaries))
  7357  }
  7358  
  7359  // Test that allocation health can't be set for an alloc with mismatched
  7360  // deployment ids
  7361  func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_MismatchDeployment(t *testing.T) {
  7362  	t.Parallel()
  7363  
  7364  	state := testStateStore(t)
  7365  
  7366  	// Insert two  deployment
  7367  	d1 := mock.Deployment()
  7368  	d2 := mock.Deployment()
  7369  	if err := state.UpsertDeployment(1, d1); err != nil {
  7370  		t.Fatalf("bad: %v", err)
  7371  	}
  7372  	if err := state.UpsertDeployment(2, d2); err != nil {
  7373  		t.Fatalf("bad: %v", err)
  7374  	}
  7375  
  7376  	// Insert an alloc for a random deployment
  7377  	a := mock.Alloc()
  7378  	a.DeploymentID = d1.ID
  7379  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 3, []*structs.Allocation{a}); err != nil {
  7380  		t.Fatalf("bad: %v", err)
  7381  	}
  7382  
  7383  	// Set health against the terminal deployment
  7384  	req := &structs.ApplyDeploymentAllocHealthRequest{
  7385  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  7386  			DeploymentID:         d2.ID,
  7387  			HealthyAllocationIDs: []string{a.ID},
  7388  		},
  7389  	}
  7390  	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 4, req)
  7391  	if err == nil || !strings.Contains(err.Error(), "not part of deployment") {
  7392  		t.Fatalf("expected error because the alloc isn't part of the deployment: %v", err)
  7393  	}
  7394  }
  7395  
  7396  // Test that allocation health is properly set
  7397  func TestStateStore_UpsertDeploymentAllocHealth(t *testing.T) {
  7398  	t.Parallel()
  7399  
  7400  	state := testStateStore(t)
  7401  
  7402  	// Insert a deployment
  7403  	d := mock.Deployment()
  7404  	d.TaskGroups["web"].ProgressDeadline = 5 * time.Minute
  7405  	if err := state.UpsertDeployment(1, d); err != nil {
  7406  		t.Fatalf("bad: %v", err)
  7407  	}
  7408  
  7409  	// Insert two allocations
  7410  	a1 := mock.Alloc()
  7411  	a1.DeploymentID = d.ID
  7412  	a2 := mock.Alloc()
  7413  	a2.DeploymentID = d.ID
  7414  	if err := state.UpsertAllocs(structs.MsgTypeTestSetup, 2, []*structs.Allocation{a1, a2}); err != nil {
  7415  		t.Fatalf("bad: %v", err)
  7416  	}
  7417  
  7418  	// Create a job to roll back to
  7419  	j := mock.Job()
  7420  
  7421  	// Create an eval that should be upserted
  7422  	e := mock.Eval()
  7423  
  7424  	// Create a status update for the deployment
  7425  	status, desc := structs.DeploymentStatusFailed, "foo"
  7426  	u := &structs.DeploymentStatusUpdate{
  7427  		DeploymentID:      d.ID,
  7428  		Status:            status,
  7429  		StatusDescription: desc,
  7430  	}
  7431  
  7432  	// Capture the time for the update
  7433  	ts := time.Now()
  7434  
  7435  	// Set health against the deployment
  7436  	req := &structs.ApplyDeploymentAllocHealthRequest{
  7437  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  7438  			DeploymentID:           d.ID,
  7439  			HealthyAllocationIDs:   []string{a1.ID},
  7440  			UnhealthyAllocationIDs: []string{a2.ID},
  7441  		},
  7442  		Job:              j,
  7443  		Eval:             e,
  7444  		DeploymentUpdate: u,
  7445  		Timestamp:        ts,
  7446  	}
  7447  	err := state.UpdateDeploymentAllocHealth(structs.MsgTypeTestSetup, 3, req)
  7448  	if err != nil {
  7449  		t.Fatalf("bad: %v", err)
  7450  	}
  7451  
  7452  	// Check that the status was updated properly
  7453  	ws := memdb.NewWatchSet()
  7454  	dout, err := state.DeploymentByID(ws, d.ID)
  7455  	if err != nil {
  7456  		t.Fatalf("bad: %v", err)
  7457  	}
  7458  	if dout.Status != status || dout.StatusDescription != desc {
  7459  		t.Fatalf("bad: %#v", dout)
  7460  	}
  7461  
  7462  	// Check that the evaluation was created
  7463  	eout, _ := state.EvalByID(ws, e.ID)
  7464  	if err != nil {
  7465  		t.Fatalf("bad: %v", err)
  7466  	}
  7467  	if eout == nil {
  7468  		t.Fatalf("bad: %#v", eout)
  7469  	}
  7470  
  7471  	// Check that the job was created
  7472  	jout, _ := state.JobByID(ws, j.Namespace, j.ID)
  7473  	if err != nil {
  7474  		t.Fatalf("bad: %v", err)
  7475  	}
  7476  	if jout == nil {
  7477  		t.Fatalf("bad: %#v", jout)
  7478  	}
  7479  
  7480  	// Check the status of the allocs
  7481  	out1, err := state.AllocByID(ws, a1.ID)
  7482  	if err != nil {
  7483  		t.Fatalf("err: %v", err)
  7484  	}
  7485  	out2, err := state.AllocByID(ws, a2.ID)
  7486  	if err != nil {
  7487  		t.Fatalf("err: %v", err)
  7488  	}
  7489  
  7490  	if !out1.DeploymentStatus.IsHealthy() {
  7491  		t.Fatalf("bad: alloc %q not healthy", out1.ID)
  7492  	}
  7493  	if !out2.DeploymentStatus.IsUnhealthy() {
  7494  		t.Fatalf("bad: alloc %q not unhealthy", out2.ID)
  7495  	}
  7496  
  7497  	if !out1.DeploymentStatus.Timestamp.Equal(ts) {
  7498  		t.Fatalf("bad: alloc %q had timestamp %v; want %v", out1.ID, out1.DeploymentStatus.Timestamp, ts)
  7499  	}
  7500  	if !out2.DeploymentStatus.Timestamp.Equal(ts) {
  7501  		t.Fatalf("bad: alloc %q had timestamp %v; want %v", out2.ID, out2.DeploymentStatus.Timestamp, ts)
  7502  	}
  7503  }
  7504  
  7505  func TestStateStore_UpsertVaultAccessors(t *testing.T) {
  7506  	t.Parallel()
  7507  
  7508  	state := testStateStore(t)
  7509  	a := mock.VaultAccessor()
  7510  	a2 := mock.VaultAccessor()
  7511  
  7512  	ws := memdb.NewWatchSet()
  7513  	if _, err := state.VaultAccessor(ws, a.Accessor); err != nil {
  7514  		t.Fatalf("err: %v", err)
  7515  	}
  7516  
  7517  	if _, err := state.VaultAccessor(ws, a2.Accessor); err != nil {
  7518  		t.Fatalf("err: %v", err)
  7519  	}
  7520  
  7521  	err := state.UpsertVaultAccessor(1000, []*structs.VaultAccessor{a, a2})
  7522  	if err != nil {
  7523  		t.Fatalf("err: %v", err)
  7524  	}
  7525  
  7526  	if !watchFired(ws) {
  7527  		t.Fatalf("bad")
  7528  	}
  7529  
  7530  	ws = memdb.NewWatchSet()
  7531  	out, err := state.VaultAccessor(ws, a.Accessor)
  7532  	if err != nil {
  7533  		t.Fatalf("err: %v", err)
  7534  	}
  7535  
  7536  	if !reflect.DeepEqual(a, out) {
  7537  		t.Fatalf("bad: %#v %#v", a, out)
  7538  	}
  7539  
  7540  	out, err = state.VaultAccessor(ws, a2.Accessor)
  7541  	if err != nil {
  7542  		t.Fatalf("err: %v", err)
  7543  	}
  7544  
  7545  	if !reflect.DeepEqual(a2, out) {
  7546  		t.Fatalf("bad: %#v %#v", a2, out)
  7547  	}
  7548  
  7549  	iter, err := state.VaultAccessors(ws)
  7550  	if err != nil {
  7551  		t.Fatalf("err: %v", err)
  7552  	}
  7553  
  7554  	count := 0
  7555  	for {
  7556  		raw := iter.Next()
  7557  		if raw == nil {
  7558  			break
  7559  		}
  7560  
  7561  		count++
  7562  		accessor := raw.(*structs.VaultAccessor)
  7563  
  7564  		if !reflect.DeepEqual(accessor, a) && !reflect.DeepEqual(accessor, a2) {
  7565  			t.Fatalf("bad: %#v", accessor)
  7566  		}
  7567  	}
  7568  
  7569  	if count != 2 {
  7570  		t.Fatalf("bad: %d", count)
  7571  	}
  7572  
  7573  	index, err := state.Index("vault_accessors")
  7574  	if err != nil {
  7575  		t.Fatalf("err: %v", err)
  7576  	}
  7577  	if index != 1000 {
  7578  		t.Fatalf("bad: %d", index)
  7579  	}
  7580  
  7581  	if watchFired(ws) {
  7582  		t.Fatalf("bad")
  7583  	}
  7584  }
  7585  
  7586  func TestStateStore_DeleteVaultAccessors(t *testing.T) {
  7587  	t.Parallel()
  7588  
  7589  	state := testStateStore(t)
  7590  	a1 := mock.VaultAccessor()
  7591  	a2 := mock.VaultAccessor()
  7592  	accessors := []*structs.VaultAccessor{a1, a2}
  7593  
  7594  	err := state.UpsertVaultAccessor(1000, accessors)
  7595  	if err != nil {
  7596  		t.Fatalf("err: %v", err)
  7597  	}
  7598  
  7599  	ws := memdb.NewWatchSet()
  7600  	if _, err := state.VaultAccessor(ws, a1.Accessor); err != nil {
  7601  		t.Fatalf("err: %v", err)
  7602  	}
  7603  
  7604  	err = state.DeleteVaultAccessors(1001, accessors)
  7605  	if err != nil {
  7606  		t.Fatalf("err: %v", err)
  7607  	}
  7608  
  7609  	if !watchFired(ws) {
  7610  		t.Fatalf("bad")
  7611  	}
  7612  
  7613  	ws = memdb.NewWatchSet()
  7614  	out, err := state.VaultAccessor(ws, a1.Accessor)
  7615  	if err != nil {
  7616  		t.Fatalf("err: %v", err)
  7617  	}
  7618  	if out != nil {
  7619  		t.Fatalf("bad: %#v %#v", a1, out)
  7620  	}
  7621  	out, err = state.VaultAccessor(ws, a2.Accessor)
  7622  	if err != nil {
  7623  		t.Fatalf("err: %v", err)
  7624  	}
  7625  	if out != nil {
  7626  		t.Fatalf("bad: %#v %#v", a2, out)
  7627  	}
  7628  
  7629  	index, err := state.Index("vault_accessors")
  7630  	if err != nil {
  7631  		t.Fatalf("err: %v", err)
  7632  	}
  7633  	if index != 1001 {
  7634  		t.Fatalf("bad: %d", index)
  7635  	}
  7636  
  7637  	if watchFired(ws) {
  7638  		t.Fatalf("bad")
  7639  	}
  7640  }
  7641  
  7642  func TestStateStore_VaultAccessorsByAlloc(t *testing.T) {
  7643  	t.Parallel()
  7644  
  7645  	state := testStateStore(t)
  7646  	alloc := mock.Alloc()
  7647  	var accessors []*structs.VaultAccessor
  7648  	var expected []*structs.VaultAccessor
  7649  
  7650  	for i := 0; i < 5; i++ {
  7651  		accessor := mock.VaultAccessor()
  7652  		accessor.AllocID = alloc.ID
  7653  		expected = append(expected, accessor)
  7654  		accessors = append(accessors, accessor)
  7655  	}
  7656  
  7657  	for i := 0; i < 10; i++ {
  7658  		accessor := mock.VaultAccessor()
  7659  		accessors = append(accessors, accessor)
  7660  	}
  7661  
  7662  	err := state.UpsertVaultAccessor(1000, accessors)
  7663  	if err != nil {
  7664  		t.Fatalf("err: %v", err)
  7665  	}
  7666  
  7667  	ws := memdb.NewWatchSet()
  7668  	out, err := state.VaultAccessorsByAlloc(ws, alloc.ID)
  7669  	if err != nil {
  7670  		t.Fatalf("err: %v", err)
  7671  	}
  7672  
  7673  	if len(expected) != len(out) {
  7674  		t.Fatalf("bad: %#v %#v", len(expected), len(out))
  7675  	}
  7676  
  7677  	index, err := state.Index("vault_accessors")
  7678  	if err != nil {
  7679  		t.Fatalf("err: %v", err)
  7680  	}
  7681  	if index != 1000 {
  7682  		t.Fatalf("bad: %d", index)
  7683  	}
  7684  
  7685  	if watchFired(ws) {
  7686  		t.Fatalf("bad")
  7687  	}
  7688  }
  7689  
  7690  func TestStateStore_VaultAccessorsByNode(t *testing.T) {
  7691  	t.Parallel()
  7692  
  7693  	state := testStateStore(t)
  7694  	node := mock.Node()
  7695  	var accessors []*structs.VaultAccessor
  7696  	var expected []*structs.VaultAccessor
  7697  
  7698  	for i := 0; i < 5; i++ {
  7699  		accessor := mock.VaultAccessor()
  7700  		accessor.NodeID = node.ID
  7701  		expected = append(expected, accessor)
  7702  		accessors = append(accessors, accessor)
  7703  	}
  7704  
  7705  	for i := 0; i < 10; i++ {
  7706  		accessor := mock.VaultAccessor()
  7707  		accessors = append(accessors, accessor)
  7708  	}
  7709  
  7710  	err := state.UpsertVaultAccessor(1000, accessors)
  7711  	if err != nil {
  7712  		t.Fatalf("err: %v", err)
  7713  	}
  7714  
  7715  	ws := memdb.NewWatchSet()
  7716  	out, err := state.VaultAccessorsByNode(ws, node.ID)
  7717  	if err != nil {
  7718  		t.Fatalf("err: %v", err)
  7719  	}
  7720  
  7721  	if len(expected) != len(out) {
  7722  		t.Fatalf("bad: %#v %#v", len(expected), len(out))
  7723  	}
  7724  
  7725  	index, err := state.Index("vault_accessors")
  7726  	if err != nil {
  7727  		t.Fatalf("err: %v", err)
  7728  	}
  7729  	if index != 1000 {
  7730  		t.Fatalf("bad: %d", index)
  7731  	}
  7732  
  7733  	if watchFired(ws) {
  7734  		t.Fatalf("bad")
  7735  	}
  7736  }
  7737  
  7738  func TestStateStore_RestoreVaultAccessor(t *testing.T) {
  7739  	t.Parallel()
  7740  
  7741  	state := testStateStore(t)
  7742  	a := mock.VaultAccessor()
  7743  
  7744  	restore, err := state.Restore()
  7745  	if err != nil {
  7746  		t.Fatalf("err: %v", err)
  7747  	}
  7748  
  7749  	err = restore.VaultAccessorRestore(a)
  7750  	if err != nil {
  7751  		t.Fatalf("err: %v", err)
  7752  	}
  7753  	require.NoError(t, restore.Commit())
  7754  
  7755  	ws := memdb.NewWatchSet()
  7756  	out, err := state.VaultAccessor(ws, a.Accessor)
  7757  	if err != nil {
  7758  		t.Fatalf("err: %v", err)
  7759  	}
  7760  
  7761  	if !reflect.DeepEqual(out, a) {
  7762  		t.Fatalf("Bad: %#v %#v", out, a)
  7763  	}
  7764  
  7765  	if watchFired(ws) {
  7766  		t.Fatalf("bad")
  7767  	}
  7768  }
  7769  
  7770  func TestStateStore_UpsertSITokenAccessors(t *testing.T) {
  7771  	t.Parallel()
  7772  	r := require.New(t)
  7773  
  7774  	state := testStateStore(t)
  7775  	a1 := mock.SITokenAccessor()
  7776  	a2 := mock.SITokenAccessor()
  7777  
  7778  	ws := memdb.NewWatchSet()
  7779  	var err error
  7780  
  7781  	_, err = state.SITokenAccessor(ws, a1.AccessorID)
  7782  	r.NoError(err)
  7783  
  7784  	_, err = state.SITokenAccessor(ws, a2.AccessorID)
  7785  	r.NoError(err)
  7786  
  7787  	err = state.UpsertSITokenAccessors(1000, []*structs.SITokenAccessor{a1, a2})
  7788  	r.NoError(err)
  7789  
  7790  	wsFired := watchFired(ws)
  7791  	r.True(wsFired)
  7792  
  7793  	noInsertWS := memdb.NewWatchSet()
  7794  	result1, err := state.SITokenAccessor(noInsertWS, a1.AccessorID)
  7795  	r.NoError(err)
  7796  	r.Equal(a1, result1)
  7797  
  7798  	result2, err := state.SITokenAccessor(noInsertWS, a2.AccessorID)
  7799  	r.NoError(err)
  7800  	r.Equal(a2, result2)
  7801  
  7802  	iter, err := state.SITokenAccessors(noInsertWS)
  7803  	r.NoError(err)
  7804  
  7805  	count := 0
  7806  	for raw := iter.Next(); raw != nil; raw = iter.Next() {
  7807  		count++
  7808  		accessor := raw.(*structs.SITokenAccessor)
  7809  		// iterator is sorted by dynamic UUID
  7810  		matches := reflect.DeepEqual(a1, accessor) || reflect.DeepEqual(a2, accessor)
  7811  		r.True(matches)
  7812  	}
  7813  	r.Equal(2, count)
  7814  
  7815  	index, err := state.Index(siTokenAccessorTable)
  7816  	r.NoError(err)
  7817  	r.Equal(uint64(1000), index)
  7818  
  7819  	noInsertWSFired := watchFired(noInsertWS)
  7820  	r.False(noInsertWSFired)
  7821  }
  7822  
  7823  func TestStateStore_DeleteSITokenAccessors(t *testing.T) {
  7824  	t.Parallel()
  7825  	r := require.New(t)
  7826  
  7827  	state := testStateStore(t)
  7828  	a1 := mock.SITokenAccessor()
  7829  	a2 := mock.SITokenAccessor()
  7830  	accessors := []*structs.SITokenAccessor{a1, a2}
  7831  	var err error
  7832  
  7833  	err = state.UpsertSITokenAccessors(1000, accessors)
  7834  	r.NoError(err)
  7835  
  7836  	ws := memdb.NewWatchSet()
  7837  	_, err = state.SITokenAccessor(ws, a1.AccessorID)
  7838  	r.NoError(err)
  7839  
  7840  	err = state.DeleteSITokenAccessors(1001, accessors)
  7841  	r.NoError(err)
  7842  
  7843  	wsFired := watchFired(ws)
  7844  	r.True(wsFired)
  7845  
  7846  	wsPostDelete := memdb.NewWatchSet()
  7847  
  7848  	result1, err := state.SITokenAccessor(wsPostDelete, a1.AccessorID)
  7849  	r.NoError(err)
  7850  	r.Nil(result1) // was deleted
  7851  
  7852  	result2, err := state.SITokenAccessor(wsPostDelete, a2.AccessorID)
  7853  	r.NoError(err)
  7854  	r.Nil(result2) // was deleted
  7855  
  7856  	index, err := state.Index(siTokenAccessorTable)
  7857  	r.NoError(err)
  7858  	r.Equal(uint64(1001), index)
  7859  
  7860  	wsPostDeleteFired := watchFired(wsPostDelete)
  7861  	r.False(wsPostDeleteFired)
  7862  }
  7863  
  7864  func TestStateStore_SITokenAccessorsByAlloc(t *testing.T) {
  7865  	t.Parallel()
  7866  	r := require.New(t)
  7867  
  7868  	state := testStateStore(t)
  7869  	alloc := mock.Alloc()
  7870  	var accessors []*structs.SITokenAccessor
  7871  	var expected []*structs.SITokenAccessor
  7872  
  7873  	for i := 0; i < 5; i++ {
  7874  		accessor := mock.SITokenAccessor()
  7875  		accessor.AllocID = alloc.ID
  7876  		expected = append(expected, accessor)
  7877  		accessors = append(accessors, accessor)
  7878  	}
  7879  
  7880  	for i := 0; i < 10; i++ {
  7881  		accessor := mock.SITokenAccessor()
  7882  		accessor.AllocID = uuid.Generate() // does not belong to alloc
  7883  		accessors = append(accessors, accessor)
  7884  	}
  7885  
  7886  	err := state.UpsertSITokenAccessors(1000, accessors)
  7887  	r.NoError(err)
  7888  
  7889  	ws := memdb.NewWatchSet()
  7890  	result, err := state.SITokenAccessorsByAlloc(ws, alloc.ID)
  7891  	r.NoError(err)
  7892  	r.ElementsMatch(expected, result)
  7893  
  7894  	index, err := state.Index(siTokenAccessorTable)
  7895  	r.NoError(err)
  7896  	r.Equal(uint64(1000), index)
  7897  
  7898  	wsFired := watchFired(ws)
  7899  	r.False(wsFired)
  7900  }
  7901  
  7902  func TestStateStore_SITokenAccessorsByNode(t *testing.T) {
  7903  	t.Parallel()
  7904  	r := require.New(t)
  7905  
  7906  	state := testStateStore(t)
  7907  	node := mock.Node()
  7908  	var accessors []*structs.SITokenAccessor
  7909  	var expected []*structs.SITokenAccessor
  7910  	var err error
  7911  
  7912  	for i := 0; i < 5; i++ {
  7913  		accessor := mock.SITokenAccessor()
  7914  		accessor.NodeID = node.ID
  7915  		expected = append(expected, accessor)
  7916  		accessors = append(accessors, accessor)
  7917  	}
  7918  
  7919  	for i := 0; i < 10; i++ {
  7920  		accessor := mock.SITokenAccessor()
  7921  		accessor.NodeID = uuid.Generate() // does not belong to node
  7922  		accessors = append(accessors, accessor)
  7923  	}
  7924  
  7925  	err = state.UpsertSITokenAccessors(1000, accessors)
  7926  	r.NoError(err)
  7927  
  7928  	ws := memdb.NewWatchSet()
  7929  	result, err := state.SITokenAccessorsByNode(ws, node.ID)
  7930  	r.NoError(err)
  7931  	r.ElementsMatch(expected, result)
  7932  
  7933  	index, err := state.Index(siTokenAccessorTable)
  7934  	r.NoError(err)
  7935  	r.Equal(uint64(1000), index)
  7936  
  7937  	wsFired := watchFired(ws)
  7938  	r.False(wsFired)
  7939  }
  7940  
  7941  func TestStateStore_RestoreSITokenAccessor(t *testing.T) {
  7942  	t.Parallel()
  7943  	r := require.New(t)
  7944  
  7945  	state := testStateStore(t)
  7946  	a1 := mock.SITokenAccessor()
  7947  
  7948  	restore, err := state.Restore()
  7949  	r.NoError(err)
  7950  
  7951  	err = restore.SITokenAccessorRestore(a1)
  7952  	r.NoError(err)
  7953  
  7954  	require.NoError(t, restore.Commit())
  7955  
  7956  	ws := memdb.NewWatchSet()
  7957  	result, err := state.SITokenAccessor(ws, a1.AccessorID)
  7958  	r.NoError(err)
  7959  	r.Equal(a1, result)
  7960  
  7961  	wsFired := watchFired(ws)
  7962  	r.False(wsFired)
  7963  }
  7964  
  7965  func TestStateStore_UpsertACLPolicy(t *testing.T) {
  7966  	t.Parallel()
  7967  
  7968  	state := testStateStore(t)
  7969  	policy := mock.ACLPolicy()
  7970  	policy2 := mock.ACLPolicy()
  7971  
  7972  	ws := memdb.NewWatchSet()
  7973  	if _, err := state.ACLPolicyByName(ws, policy.Name); err != nil {
  7974  		t.Fatalf("err: %v", err)
  7975  	}
  7976  	if _, err := state.ACLPolicyByName(ws, policy2.Name); err != nil {
  7977  		t.Fatalf("err: %v", err)
  7978  	}
  7979  
  7980  	if err := state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy, policy2}); err != nil {
  7981  		t.Fatalf("err: %v", err)
  7982  	}
  7983  	if !watchFired(ws) {
  7984  		t.Fatalf("bad")
  7985  	}
  7986  
  7987  	ws = memdb.NewWatchSet()
  7988  	out, err := state.ACLPolicyByName(ws, policy.Name)
  7989  	assert.Equal(t, nil, err)
  7990  	assert.Equal(t, policy, out)
  7991  
  7992  	out, err = state.ACLPolicyByName(ws, policy2.Name)
  7993  	assert.Equal(t, nil, err)
  7994  	assert.Equal(t, policy2, out)
  7995  
  7996  	iter, err := state.ACLPolicies(ws)
  7997  	if err != nil {
  7998  		t.Fatalf("err: %v", err)
  7999  	}
  8000  
  8001  	// Ensure we see both policies
  8002  	count := 0
  8003  	for {
  8004  		raw := iter.Next()
  8005  		if raw == nil {
  8006  			break
  8007  		}
  8008  		count++
  8009  	}
  8010  	if count != 2 {
  8011  		t.Fatalf("bad: %d", count)
  8012  	}
  8013  
  8014  	index, err := state.Index("acl_policy")
  8015  	if err != nil {
  8016  		t.Fatalf("err: %v", err)
  8017  	}
  8018  	if index != 1000 {
  8019  		t.Fatalf("bad: %d", index)
  8020  	}
  8021  
  8022  	if watchFired(ws) {
  8023  		t.Fatalf("bad")
  8024  	}
  8025  }
  8026  
  8027  func TestStateStore_DeleteACLPolicy(t *testing.T) {
  8028  	t.Parallel()
  8029  
  8030  	state := testStateStore(t)
  8031  	policy := mock.ACLPolicy()
  8032  	policy2 := mock.ACLPolicy()
  8033  
  8034  	// Create the policy
  8035  	if err := state.UpsertACLPolicies(structs.MsgTypeTestSetup, 1000, []*structs.ACLPolicy{policy, policy2}); err != nil {
  8036  		t.Fatalf("err: %v", err)
  8037  	}
  8038  
  8039  	// Create a watcher
  8040  	ws := memdb.NewWatchSet()
  8041  	if _, err := state.ACLPolicyByName(ws, policy.Name); err != nil {
  8042  		t.Fatalf("err: %v", err)
  8043  	}
  8044  
  8045  	// Delete the policy
  8046  	if err := state.DeleteACLPolicies(structs.MsgTypeTestSetup, 1001, []string{policy.Name, policy2.Name}); err != nil {
  8047  		t.Fatalf("err: %v", err)
  8048  	}
  8049  
  8050  	// Ensure watching triggered
  8051  	if !watchFired(ws) {
  8052  		t.Fatalf("bad")
  8053  	}
  8054  
  8055  	// Ensure we don't get the object back
  8056  	ws = memdb.NewWatchSet()
  8057  	out, err := state.ACLPolicyByName(ws, policy.Name)
  8058  	assert.Equal(t, nil, err)
  8059  	if out != nil {
  8060  		t.Fatalf("bad: %#v", out)
  8061  	}
  8062  
  8063  	iter, err := state.ACLPolicies(ws)
  8064  	if err != nil {
  8065  		t.Fatalf("err: %v", err)
  8066  	}
  8067  
  8068  	// Ensure we see neither policy
  8069  	count := 0
  8070  	for {
  8071  		raw := iter.Next()
  8072  		if raw == nil {
  8073  			break
  8074  		}
  8075  		count++
  8076  	}
  8077  	if count != 0 {
  8078  		t.Fatalf("bad: %d", count)
  8079  	}
  8080  
  8081  	index, err := state.Index("acl_policy")
  8082  	if err != nil {
  8083  		t.Fatalf("err: %v", err)
  8084  	}
  8085  	if index != 1001 {
  8086  		t.Fatalf("bad: %d", index)
  8087  	}
  8088  
  8089  	if watchFired(ws) {
  8090  		t.Fatalf("bad")
  8091  	}
  8092  }
  8093  
  8094  func TestStateStore_ACLPolicyByNamePrefix(t *testing.T) {
  8095  	t.Parallel()
  8096  
  8097  	state := testStateStore(t)
  8098  	names := []string{
  8099  		"foo",
  8100  		"bar",
  8101  		"foobar",
  8102  		"foozip",
  8103  		"zip",
  8104  	}
  8105  
  8106  	// Create the policies
  8107  	var baseIndex uint64 = 1000
  8108  	for _, name := range names {
  8109  		p := mock.ACLPolicy()
  8110  		p.Name = name
  8111  		if err := state.UpsertACLPolicies(structs.MsgTypeTestSetup, baseIndex, []*structs.ACLPolicy{p}); err != nil {
  8112  			t.Fatalf("err: %v", err)
  8113  		}
  8114  		baseIndex++
  8115  	}
  8116  
  8117  	// Scan by prefix
  8118  	iter, err := state.ACLPolicyByNamePrefix(nil, "foo")
  8119  	if err != nil {
  8120  		t.Fatalf("err: %v", err)
  8121  	}
  8122  
  8123  	// Ensure we see both policies
  8124  	count := 0
  8125  	out := []string{}
  8126  	for {
  8127  		raw := iter.Next()
  8128  		if raw == nil {
  8129  			break
  8130  		}
  8131  		count++
  8132  		out = append(out, raw.(*structs.ACLPolicy).Name)
  8133  	}
  8134  	if count != 3 {
  8135  		t.Fatalf("bad: %d %v", count, out)
  8136  	}
  8137  	sort.Strings(out)
  8138  
  8139  	expect := []string{"foo", "foobar", "foozip"}
  8140  	assert.Equal(t, expect, out)
  8141  }
  8142  
  8143  func TestStateStore_BootstrapACLTokens(t *testing.T) {
  8144  	t.Parallel()
  8145  
  8146  	state := testStateStore(t)
  8147  	tk1 := mock.ACLToken()
  8148  	tk2 := mock.ACLToken()
  8149  
  8150  	ok, resetIdx, err := state.CanBootstrapACLToken()
  8151  	assert.Nil(t, err)
  8152  	assert.Equal(t, true, ok)
  8153  	assert.EqualValues(t, 0, resetIdx)
  8154  
  8155  	if err := state.BootstrapACLTokens(structs.MsgTypeTestSetup, 1000, 0, tk1); err != nil {
  8156  		t.Fatalf("err: %v", err)
  8157  	}
  8158  
  8159  	out, err := state.ACLTokenByAccessorID(nil, tk1.AccessorID)
  8160  	assert.Equal(t, nil, err)
  8161  	assert.Equal(t, tk1, out)
  8162  
  8163  	ok, resetIdx, err = state.CanBootstrapACLToken()
  8164  	assert.Nil(t, err)
  8165  	assert.Equal(t, false, ok)
  8166  	assert.EqualValues(t, 1000, resetIdx)
  8167  
  8168  	if err := state.BootstrapACLTokens(structs.MsgTypeTestSetup, 1001, 0, tk2); err == nil {
  8169  		t.Fatalf("expected error")
  8170  	}
  8171  
  8172  	iter, err := state.ACLTokens(nil)
  8173  	if err != nil {
  8174  		t.Fatalf("err: %v", err)
  8175  	}
  8176  
  8177  	// Ensure we see both policies
  8178  	count := 0
  8179  	for {
  8180  		raw := iter.Next()
  8181  		if raw == nil {
  8182  			break
  8183  		}
  8184  		count++
  8185  	}
  8186  	if count != 1 {
  8187  		t.Fatalf("bad: %d", count)
  8188  	}
  8189  
  8190  	index, err := state.Index("acl_token")
  8191  	if err != nil {
  8192  		t.Fatalf("err: %v", err)
  8193  	}
  8194  	if index != 1000 {
  8195  		t.Fatalf("bad: %d", index)
  8196  	}
  8197  	index, err = state.Index("acl_token_bootstrap")
  8198  	if err != nil {
  8199  		t.Fatalf("err: %v", err)
  8200  	}
  8201  	if index != 1000 {
  8202  		t.Fatalf("bad: %d", index)
  8203  	}
  8204  
  8205  	// Should allow bootstrap with reset index
  8206  	if err := state.BootstrapACLTokens(structs.MsgTypeTestSetup, 1001, 1000, tk2); err != nil {
  8207  		t.Fatalf("err %v", err)
  8208  	}
  8209  
  8210  	// Check we've modified the index
  8211  	index, err = state.Index("acl_token")
  8212  	if err != nil {
  8213  		t.Fatalf("err: %v", err)
  8214  	}
  8215  	if index != 1001 {
  8216  		t.Fatalf("bad: %d", index)
  8217  	}
  8218  	index, err = state.Index("acl_token_bootstrap")
  8219  	if err != nil {
  8220  		t.Fatalf("err: %v", err)
  8221  	}
  8222  	if index != 1001 {
  8223  		t.Fatalf("bad: %d", index)
  8224  	}
  8225  }
  8226  
  8227  func TestStateStore_UpsertACLTokens(t *testing.T) {
  8228  	t.Parallel()
  8229  
  8230  	state := testStateStore(t)
  8231  	tk1 := mock.ACLToken()
  8232  	tk2 := mock.ACLToken()
  8233  
  8234  	ws := memdb.NewWatchSet()
  8235  	if _, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID); err != nil {
  8236  		t.Fatalf("err: %v", err)
  8237  	}
  8238  	if _, err := state.ACLTokenByAccessorID(ws, tk2.AccessorID); err != nil {
  8239  		t.Fatalf("err: %v", err)
  8240  	}
  8241  
  8242  	if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{tk1, tk2}); err != nil {
  8243  		t.Fatalf("err: %v", err)
  8244  	}
  8245  	if !watchFired(ws) {
  8246  		t.Fatalf("bad")
  8247  	}
  8248  
  8249  	ws = memdb.NewWatchSet()
  8250  	out, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID)
  8251  	assert.Equal(t, nil, err)
  8252  	assert.Equal(t, tk1, out)
  8253  
  8254  	out, err = state.ACLTokenByAccessorID(ws, tk2.AccessorID)
  8255  	assert.Equal(t, nil, err)
  8256  	assert.Equal(t, tk2, out)
  8257  
  8258  	out, err = state.ACLTokenBySecretID(ws, tk1.SecretID)
  8259  	assert.Equal(t, nil, err)
  8260  	assert.Equal(t, tk1, out)
  8261  
  8262  	out, err = state.ACLTokenBySecretID(ws, tk2.SecretID)
  8263  	assert.Equal(t, nil, err)
  8264  	assert.Equal(t, tk2, out)
  8265  
  8266  	iter, err := state.ACLTokens(ws)
  8267  	if err != nil {
  8268  		t.Fatalf("err: %v", err)
  8269  	}
  8270  
  8271  	// Ensure we see both policies
  8272  	count := 0
  8273  	for {
  8274  		raw := iter.Next()
  8275  		if raw == nil {
  8276  			break
  8277  		}
  8278  		count++
  8279  	}
  8280  	if count != 2 {
  8281  		t.Fatalf("bad: %d", count)
  8282  	}
  8283  
  8284  	index, err := state.Index("acl_token")
  8285  	if err != nil {
  8286  		t.Fatalf("err: %v", err)
  8287  	}
  8288  	if index != 1000 {
  8289  		t.Fatalf("bad: %d", index)
  8290  	}
  8291  
  8292  	if watchFired(ws) {
  8293  		t.Fatalf("bad")
  8294  	}
  8295  }
  8296  
  8297  func TestStateStore_DeleteACLTokens(t *testing.T) {
  8298  	t.Parallel()
  8299  
  8300  	state := testStateStore(t)
  8301  	tk1 := mock.ACLToken()
  8302  	tk2 := mock.ACLToken()
  8303  
  8304  	// Create the tokens
  8305  	if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{tk1, tk2}); err != nil {
  8306  		t.Fatalf("err: %v", err)
  8307  	}
  8308  
  8309  	// Create a watcher
  8310  	ws := memdb.NewWatchSet()
  8311  	if _, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID); err != nil {
  8312  		t.Fatalf("err: %v", err)
  8313  	}
  8314  
  8315  	// Delete the token
  8316  	if err := state.DeleteACLTokens(structs.MsgTypeTestSetup, 1001, []string{tk1.AccessorID, tk2.AccessorID}); err != nil {
  8317  		t.Fatalf("err: %v", err)
  8318  	}
  8319  
  8320  	// Ensure watching triggered
  8321  	if !watchFired(ws) {
  8322  		t.Fatalf("bad")
  8323  	}
  8324  
  8325  	// Ensure we don't get the object back
  8326  	ws = memdb.NewWatchSet()
  8327  	out, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID)
  8328  	assert.Equal(t, nil, err)
  8329  	if out != nil {
  8330  		t.Fatalf("bad: %#v", out)
  8331  	}
  8332  
  8333  	iter, err := state.ACLTokens(ws)
  8334  	if err != nil {
  8335  		t.Fatalf("err: %v", err)
  8336  	}
  8337  
  8338  	// Ensure we see both policies
  8339  	count := 0
  8340  	for {
  8341  		raw := iter.Next()
  8342  		if raw == nil {
  8343  			break
  8344  		}
  8345  		count++
  8346  	}
  8347  	if count != 0 {
  8348  		t.Fatalf("bad: %d", count)
  8349  	}
  8350  
  8351  	index, err := state.Index("acl_token")
  8352  	if err != nil {
  8353  		t.Fatalf("err: %v", err)
  8354  	}
  8355  	if index != 1001 {
  8356  		t.Fatalf("bad: %d", index)
  8357  	}
  8358  
  8359  	if watchFired(ws) {
  8360  		t.Fatalf("bad")
  8361  	}
  8362  }
  8363  
  8364  func TestStateStore_ACLTokenByAccessorIDPrefix(t *testing.T) {
  8365  	t.Parallel()
  8366  
  8367  	state := testStateStore(t)
  8368  	prefixes := []string{
  8369  		"aaaa",
  8370  		"aabb",
  8371  		"bbbb",
  8372  		"bbcc",
  8373  		"ffff",
  8374  	}
  8375  
  8376  	// Create the tokens
  8377  	var baseIndex uint64 = 1000
  8378  	for _, prefix := range prefixes {
  8379  		tk := mock.ACLToken()
  8380  		tk.AccessorID = prefix + tk.AccessorID[4:]
  8381  		if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, baseIndex, []*structs.ACLToken{tk}); err != nil {
  8382  			t.Fatalf("err: %v", err)
  8383  		}
  8384  		baseIndex++
  8385  	}
  8386  
  8387  	// Scan by prefix
  8388  	iter, err := state.ACLTokenByAccessorIDPrefix(nil, "aa")
  8389  	if err != nil {
  8390  		t.Fatalf("err: %v", err)
  8391  	}
  8392  
  8393  	// Ensure we see both tokens
  8394  	count := 0
  8395  	out := []string{}
  8396  	for {
  8397  		raw := iter.Next()
  8398  		if raw == nil {
  8399  			break
  8400  		}
  8401  		count++
  8402  		out = append(out, raw.(*structs.ACLToken).AccessorID[:4])
  8403  	}
  8404  	if count != 2 {
  8405  		t.Fatalf("bad: %d %v", count, out)
  8406  	}
  8407  	sort.Strings(out)
  8408  
  8409  	expect := []string{"aaaa", "aabb"}
  8410  	assert.Equal(t, expect, out)
  8411  }
  8412  
  8413  func TestStateStore_RestoreACLPolicy(t *testing.T) {
  8414  	t.Parallel()
  8415  
  8416  	state := testStateStore(t)
  8417  	policy := mock.ACLPolicy()
  8418  
  8419  	restore, err := state.Restore()
  8420  	if err != nil {
  8421  		t.Fatalf("err: %v", err)
  8422  	}
  8423  
  8424  	err = restore.ACLPolicyRestore(policy)
  8425  	if err != nil {
  8426  		t.Fatalf("err: %v", err)
  8427  	}
  8428  	require.NoError(t, restore.Commit())
  8429  
  8430  	ws := memdb.NewWatchSet()
  8431  	out, err := state.ACLPolicyByName(ws, policy.Name)
  8432  	if err != nil {
  8433  		t.Fatalf("err: %v", err)
  8434  	}
  8435  	assert.Equal(t, policy, out)
  8436  }
  8437  
  8438  func TestStateStore_ACLTokensByGlobal(t *testing.T) {
  8439  	t.Parallel()
  8440  
  8441  	state := testStateStore(t)
  8442  	tk1 := mock.ACLToken()
  8443  	tk2 := mock.ACLToken()
  8444  	tk3 := mock.ACLToken()
  8445  	tk4 := mock.ACLToken()
  8446  	tk3.Global = true
  8447  
  8448  	if err := state.UpsertACLTokens(structs.MsgTypeTestSetup, 1000, []*structs.ACLToken{tk1, tk2, tk3, tk4}); err != nil {
  8449  		t.Fatalf("err: %v", err)
  8450  	}
  8451  
  8452  	iter, err := state.ACLTokensByGlobal(nil, true)
  8453  	if err != nil {
  8454  		t.Fatalf("err: %v", err)
  8455  	}
  8456  
  8457  	// Ensure we see the one global policies
  8458  	count := 0
  8459  	for {
  8460  		raw := iter.Next()
  8461  		if raw == nil {
  8462  			break
  8463  		}
  8464  		count++
  8465  	}
  8466  	if count != 1 {
  8467  		t.Fatalf("bad: %d", count)
  8468  	}
  8469  }
  8470  
  8471  func TestStateStore_RestoreACLToken(t *testing.T) {
  8472  	t.Parallel()
  8473  
  8474  	state := testStateStore(t)
  8475  	token := mock.ACLToken()
  8476  
  8477  	restore, err := state.Restore()
  8478  	if err != nil {
  8479  		t.Fatalf("err: %v", err)
  8480  	}
  8481  
  8482  	err = restore.ACLTokenRestore(token)
  8483  	if err != nil {
  8484  		t.Fatalf("err: %v", err)
  8485  	}
  8486  	require.NoError(t, restore.Commit())
  8487  
  8488  	ws := memdb.NewWatchSet()
  8489  	out, err := state.ACLTokenByAccessorID(ws, token.AccessorID)
  8490  	if err != nil {
  8491  		t.Fatalf("err: %v", err)
  8492  	}
  8493  	assert.Equal(t, token, out)
  8494  }
  8495  
  8496  func TestStateStore_SchedulerConfig(t *testing.T) {
  8497  	t.Parallel()
  8498  
  8499  	state := testStateStore(t)
  8500  	schedConfig := &structs.SchedulerConfiguration{
  8501  		PreemptionConfig: structs.PreemptionConfig{
  8502  			SystemSchedulerEnabled: false,
  8503  		},
  8504  		CreateIndex: 100,
  8505  		ModifyIndex: 200,
  8506  	}
  8507  
  8508  	require := require.New(t)
  8509  	restore, err := state.Restore()
  8510  	require.Nil(err)
  8511  
  8512  	err = restore.SchedulerConfigRestore(schedConfig)
  8513  	require.Nil(err)
  8514  
  8515  	require.NoError(restore.Commit())
  8516  
  8517  	modIndex, out, err := state.SchedulerConfig()
  8518  	require.Nil(err)
  8519  	require.Equal(schedConfig.ModifyIndex, modIndex)
  8520  
  8521  	require.Equal(schedConfig, out)
  8522  }
  8523  
  8524  func TestStateStore_ClusterMetadata(t *testing.T) {
  8525  	require := require.New(t)
  8526  
  8527  	state := testStateStore(t)
  8528  	clusterID := "12345678-1234-1234-1234-1234567890"
  8529  	now := time.Now().UnixNano()
  8530  	meta := &structs.ClusterMetadata{ClusterID: clusterID, CreateTime: now}
  8531  
  8532  	err := state.ClusterSetMetadata(100, meta)
  8533  	require.NoError(err)
  8534  
  8535  	result, err := state.ClusterMetadata(nil)
  8536  	require.NoError(err)
  8537  	require.Equal(clusterID, result.ClusterID)
  8538  	require.Equal(now, result.CreateTime)
  8539  }
  8540  
  8541  func TestStateStore_ClusterMetadataRestore(t *testing.T) {
  8542  	require := require.New(t)
  8543  
  8544  	state := testStateStore(t)
  8545  	clusterID := "12345678-1234-1234-1234-1234567890"
  8546  	now := time.Now().UnixNano()
  8547  	meta := &structs.ClusterMetadata{ClusterID: clusterID, CreateTime: now}
  8548  
  8549  	restore, err := state.Restore()
  8550  	require.NoError(err)
  8551  
  8552  	err = restore.ClusterMetadataRestore(meta)
  8553  	require.NoError(err)
  8554  
  8555  	require.NoError(restore.Commit())
  8556  
  8557  	out, err := state.ClusterMetadata(nil)
  8558  	require.NoError(err)
  8559  	require.Equal(clusterID, out.ClusterID)
  8560  	require.Equal(now, out.CreateTime)
  8561  }
  8562  
  8563  func TestStateStore_RestoreScalingPolicy(t *testing.T) {
  8564  	t.Parallel()
  8565  	require := require.New(t)
  8566  
  8567  	state := testStateStore(t)
  8568  	scalingPolicy := mock.ScalingPolicy()
  8569  
  8570  	restore, err := state.Restore()
  8571  	require.NoError(err)
  8572  
  8573  	err = restore.ScalingPolicyRestore(scalingPolicy)
  8574  	require.NoError(err)
  8575  	require.NoError(restore.Commit())
  8576  
  8577  	ws := memdb.NewWatchSet()
  8578  	out, err := state.ScalingPolicyByID(ws, scalingPolicy.ID)
  8579  	require.NoError(err)
  8580  	require.EqualValues(out, scalingPolicy)
  8581  }
  8582  
  8583  func TestStateStore_UpsertScalingPolicy(t *testing.T) {
  8584  	t.Parallel()
  8585  	require := require.New(t)
  8586  
  8587  	state := testStateStore(t)
  8588  	policy := mock.ScalingPolicy()
  8589  	policy2 := mock.ScalingPolicy()
  8590  
  8591  	wsAll := memdb.NewWatchSet()
  8592  	all, err := state.ScalingPolicies(wsAll)
  8593  	require.NoError(err)
  8594  	require.Nil(all.Next())
  8595  
  8596  	ws := memdb.NewWatchSet()
  8597  	out, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  8598  	require.NoError(err)
  8599  	require.Nil(out)
  8600  
  8601  	out, err = state.ScalingPolicyByTargetAndType(ws, policy2.Target, policy2.Type)
  8602  	require.NoError(err)
  8603  	require.Nil(out)
  8604  
  8605  	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy, policy2})
  8606  	require.NoError(err)
  8607  	require.True(watchFired(ws))
  8608  	require.True(watchFired(wsAll))
  8609  
  8610  	ws = memdb.NewWatchSet()
  8611  	out, err = state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  8612  	require.NoError(err)
  8613  	require.Equal(policy, out)
  8614  
  8615  	out, err = state.ScalingPolicyByTargetAndType(ws, policy2.Target, policy2.Type)
  8616  	require.NoError(err)
  8617  	require.Equal(policy2, out)
  8618  
  8619  	// Ensure we see both policies
  8620  	countPolicies := func() (n int, err error) {
  8621  		iter, err := state.ScalingPolicies(ws)
  8622  		if err != nil {
  8623  			return
  8624  		}
  8625  
  8626  		for raw := iter.Next(); raw != nil; raw = iter.Next() {
  8627  			n++
  8628  		}
  8629  		return
  8630  	}
  8631  
  8632  	count, err := countPolicies()
  8633  	require.NoError(err)
  8634  	require.Equal(2, count)
  8635  
  8636  	index, err := state.Index("scaling_policy")
  8637  	require.NoError(err)
  8638  	require.True(1000 == index)
  8639  	require.False(watchFired(ws))
  8640  
  8641  	// Check that we can add policy with same target but different type
  8642  	policy3 := mock.ScalingPolicy()
  8643  	for k, v := range policy2.Target {
  8644  		policy3.Target[k] = v
  8645  	}
  8646  
  8647  	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy3})
  8648  	require.NoError(err)
  8649  
  8650  	// Ensure we see both policies, since target didn't change
  8651  	count, err = countPolicies()
  8652  	require.NoError(err)
  8653  	require.Equal(2, count)
  8654  
  8655  	// Change type and check if we see 3
  8656  	policy3.Type = "other-type"
  8657  
  8658  	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy3})
  8659  	require.NoError(err)
  8660  
  8661  	count, err = countPolicies()
  8662  	require.NoError(err)
  8663  	require.Equal(3, count)
  8664  }
  8665  
  8666  func TestStateStore_UpsertScalingPolicy_Namespace(t *testing.T) {
  8667  	t.Parallel()
  8668  	require := require.New(t)
  8669  
  8670  	otherNamespace := "not-default-namespace"
  8671  	state := testStateStore(t)
  8672  	policy := mock.ScalingPolicy()
  8673  	policy2 := mock.ScalingPolicy()
  8674  	policy2.Target[structs.ScalingTargetNamespace] = otherNamespace
  8675  
  8676  	ws1 := memdb.NewWatchSet()
  8677  	iter, err := state.ScalingPoliciesByNamespace(ws1, structs.DefaultNamespace, "")
  8678  	require.NoError(err)
  8679  	require.Nil(iter.Next())
  8680  
  8681  	ws2 := memdb.NewWatchSet()
  8682  	iter, err = state.ScalingPoliciesByNamespace(ws2, otherNamespace, "")
  8683  	require.NoError(err)
  8684  	require.Nil(iter.Next())
  8685  
  8686  	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy, policy2})
  8687  	require.NoError(err)
  8688  	require.True(watchFired(ws1))
  8689  	require.True(watchFired(ws2))
  8690  
  8691  	iter, err = state.ScalingPoliciesByNamespace(nil, structs.DefaultNamespace, "")
  8692  	require.NoError(err)
  8693  	policiesInDefaultNamespace := []string{}
  8694  	for {
  8695  		raw := iter.Next()
  8696  		if raw == nil {
  8697  			break
  8698  		}
  8699  		policiesInDefaultNamespace = append(policiesInDefaultNamespace, raw.(*structs.ScalingPolicy).ID)
  8700  	}
  8701  	require.ElementsMatch([]string{policy.ID}, policiesInDefaultNamespace)
  8702  
  8703  	iter, err = state.ScalingPoliciesByNamespace(nil, otherNamespace, "")
  8704  	require.NoError(err)
  8705  	policiesInOtherNamespace := []string{}
  8706  	for {
  8707  		raw := iter.Next()
  8708  		if raw == nil {
  8709  			break
  8710  		}
  8711  		policiesInOtherNamespace = append(policiesInOtherNamespace, raw.(*structs.ScalingPolicy).ID)
  8712  	}
  8713  	require.ElementsMatch([]string{policy2.ID}, policiesInOtherNamespace)
  8714  }
  8715  
  8716  func TestStateStore_UpsertScalingPolicy_Namespace_PrefixBug(t *testing.T) {
  8717  	t.Parallel()
  8718  	require := require.New(t)
  8719  
  8720  	ns1 := "name"
  8721  	ns2 := "name2" // matches prefix "name"
  8722  	state := testStateStore(t)
  8723  	policy1 := mock.ScalingPolicy()
  8724  	policy1.Target[structs.ScalingTargetNamespace] = ns1
  8725  	policy2 := mock.ScalingPolicy()
  8726  	policy2.Target[structs.ScalingTargetNamespace] = ns2
  8727  
  8728  	ws1 := memdb.NewWatchSet()
  8729  	iter, err := state.ScalingPoliciesByNamespace(ws1, ns1, "")
  8730  	require.NoError(err)
  8731  	require.Nil(iter.Next())
  8732  
  8733  	ws2 := memdb.NewWatchSet()
  8734  	iter, err = state.ScalingPoliciesByNamespace(ws2, ns2, "")
  8735  	require.NoError(err)
  8736  	require.Nil(iter.Next())
  8737  
  8738  	err = state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy1, policy2})
  8739  	require.NoError(err)
  8740  	require.True(watchFired(ws1))
  8741  	require.True(watchFired(ws2))
  8742  
  8743  	iter, err = state.ScalingPoliciesByNamespace(nil, ns1, "")
  8744  	require.NoError(err)
  8745  	policiesInNS1 := []string{}
  8746  	for {
  8747  		raw := iter.Next()
  8748  		if raw == nil {
  8749  			break
  8750  		}
  8751  		policiesInNS1 = append(policiesInNS1, raw.(*structs.ScalingPolicy).ID)
  8752  	}
  8753  	require.ElementsMatch([]string{policy1.ID}, policiesInNS1)
  8754  
  8755  	iter, err = state.ScalingPoliciesByNamespace(nil, ns2, "")
  8756  	require.NoError(err)
  8757  	policiesInNS2 := []string{}
  8758  	for {
  8759  		raw := iter.Next()
  8760  		if raw == nil {
  8761  			break
  8762  		}
  8763  		policiesInNS2 = append(policiesInNS2, raw.(*structs.ScalingPolicy).ID)
  8764  	}
  8765  	require.ElementsMatch([]string{policy2.ID}, policiesInNS2)
  8766  }
  8767  
  8768  // Scaling Policy IDs are generated randomly during Job.Register
  8769  // Subsequent updates of the job should preserve the ID for the scaling policy
  8770  // associated with a given target.
  8771  func TestStateStore_UpsertJob_PreserveScalingPolicyIDsAndIndex(t *testing.T) {
  8772  	t.Parallel()
  8773  
  8774  	require := require.New(t)
  8775  
  8776  	state := testStateStore(t)
  8777  	job, policy := mock.JobWithScalingPolicy()
  8778  
  8779  	var newIndex uint64 = 1000
  8780  	err := state.UpsertJob(structs.MsgTypeTestSetup, newIndex, job)
  8781  	require.NoError(err)
  8782  
  8783  	ws := memdb.NewWatchSet()
  8784  	p1, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  8785  	require.NoError(err)
  8786  	require.NotNil(p1)
  8787  	require.Equal(newIndex, p1.CreateIndex)
  8788  	require.Equal(newIndex, p1.ModifyIndex)
  8789  
  8790  	index, err := state.Index("scaling_policy")
  8791  	require.NoError(err)
  8792  	require.Equal(newIndex, index)
  8793  	require.NotEmpty(p1.ID)
  8794  
  8795  	// update the job
  8796  	job.Meta["new-meta"] = "new-value"
  8797  	newIndex += 100
  8798  	err = state.UpsertJob(structs.MsgTypeTestSetup, newIndex, job)
  8799  	require.NoError(err)
  8800  	require.False(watchFired(ws), "watch should not have fired")
  8801  
  8802  	p2, err := state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
  8803  	require.NoError(err)
  8804  	require.NotNil(p2)
  8805  	require.Equal(p1.ID, p2.ID, "ID should not have changed")
  8806  	require.Equal(p1.CreateIndex, p2.CreateIndex)
  8807  	require.Equal(p1.ModifyIndex, p2.ModifyIndex)
  8808  
  8809  	index, err = state.Index("scaling_policy")
  8810  	require.NoError(err)
  8811  	require.Equal(index, p1.CreateIndex, "table index should not have changed")
  8812  }
  8813  
  8814  // Updating the scaling policy for a job should update the index table and fire the watch.
  8815  // This test is the converse of TestStateStore_UpsertJob_PreserveScalingPolicyIDsAndIndex
  8816  func TestStateStore_UpsertJob_UpdateScalingPolicy(t *testing.T) {
  8817  	t.Parallel()
  8818  
  8819  	require := require.New(t)
  8820  
  8821  	state := testStateStore(t)
  8822  	job, policy := mock.JobWithScalingPolicy()
  8823  
  8824  	var oldIndex uint64 = 1000
  8825  	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, oldIndex, job))
  8826  
  8827  	ws := memdb.NewWatchSet()
  8828  	p1, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  8829  	require.NoError(err)
  8830  	require.NotNil(p1)
  8831  	require.Equal(oldIndex, p1.CreateIndex)
  8832  	require.Equal(oldIndex, p1.ModifyIndex)
  8833  	prevId := p1.ID
  8834  
  8835  	index, err := state.Index("scaling_policy")
  8836  	require.NoError(err)
  8837  	require.Equal(oldIndex, index)
  8838  	require.NotEmpty(p1.ID)
  8839  
  8840  	// update the job with the updated scaling policy; make sure to use a different object
  8841  	newPolicy := p1.Copy()
  8842  	newPolicy.Policy["new-field"] = "new-value"
  8843  	job.TaskGroups[0].Scaling = newPolicy
  8844  	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, oldIndex+100, job))
  8845  	require.True(watchFired(ws), "watch should have fired")
  8846  
  8847  	p2, err := state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
  8848  	require.NoError(err)
  8849  	require.NotNil(p2)
  8850  	require.Equal(p2.Policy["new-field"], "new-value")
  8851  	require.Equal(prevId, p2.ID, "ID should not have changed")
  8852  	require.Equal(oldIndex, p2.CreateIndex)
  8853  	require.Greater(p2.ModifyIndex, oldIndex, "ModifyIndex should have advanced")
  8854  
  8855  	index, err = state.Index("scaling_policy")
  8856  	require.NoError(err)
  8857  	require.Greater(index, oldIndex, "table index should have advanced")
  8858  }
  8859  
  8860  func TestStateStore_DeleteScalingPolicies(t *testing.T) {
  8861  	t.Parallel()
  8862  
  8863  	require := require.New(t)
  8864  
  8865  	state := testStateStore(t)
  8866  	policy := mock.ScalingPolicy()
  8867  	policy2 := mock.ScalingPolicy()
  8868  
  8869  	// Create the policy
  8870  	err := state.UpsertScalingPolicies(1000, []*structs.ScalingPolicy{policy, policy2})
  8871  	require.NoError(err)
  8872  
  8873  	// Create a watcher
  8874  	ws := memdb.NewWatchSet()
  8875  	_, err = state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  8876  	require.NoError(err)
  8877  
  8878  	// Delete the policy
  8879  	err = state.DeleteScalingPolicies(1001, []string{policy.ID, policy2.ID})
  8880  	require.NoError(err)
  8881  
  8882  	// Ensure watching triggered
  8883  	require.True(watchFired(ws))
  8884  
  8885  	// Ensure we don't get the objects back
  8886  	ws = memdb.NewWatchSet()
  8887  	out, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  8888  	require.NoError(err)
  8889  	require.Nil(out)
  8890  
  8891  	ws = memdb.NewWatchSet()
  8892  	out, err = state.ScalingPolicyByTargetAndType(ws, policy2.Target, policy2.Type)
  8893  	require.NoError(err)
  8894  	require.Nil(out)
  8895  
  8896  	// Ensure we see both policies
  8897  	iter, err := state.ScalingPoliciesByNamespace(ws, policy.Target[structs.ScalingTargetNamespace], "")
  8898  	require.NoError(err)
  8899  	count := 0
  8900  	for {
  8901  		raw := iter.Next()
  8902  		if raw == nil {
  8903  			break
  8904  		}
  8905  		count++
  8906  	}
  8907  	require.Equal(0, count)
  8908  
  8909  	index, err := state.Index("scaling_policy")
  8910  	require.NoError(err)
  8911  	require.True(1001 == index)
  8912  	require.False(watchFired(ws))
  8913  }
  8914  
  8915  func TestStateStore_StopJob_DeleteScalingPolicies(t *testing.T) {
  8916  	t.Parallel()
  8917  
  8918  	require := require.New(t)
  8919  
  8920  	state := testStateStore(t)
  8921  
  8922  	job := mock.Job()
  8923  
  8924  	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  8925  	require.NoError(err)
  8926  
  8927  	policy := mock.ScalingPolicy()
  8928  	policy.Target[structs.ScalingTargetJob] = job.ID
  8929  	err = state.UpsertScalingPolicies(1100, []*structs.ScalingPolicy{policy})
  8930  	require.NoError(err)
  8931  
  8932  	// Ensure the scaling policy is present and start some watches
  8933  	wsGet := memdb.NewWatchSet()
  8934  	out, err := state.ScalingPolicyByTargetAndType(wsGet, policy.Target, policy.Type)
  8935  	require.NoError(err)
  8936  	require.NotNil(out)
  8937  	wsList := memdb.NewWatchSet()
  8938  	_, err = state.ScalingPolicies(wsList)
  8939  	require.NoError(err)
  8940  
  8941  	// Stop the job
  8942  	job, err = state.JobByID(nil, job.Namespace, job.ID)
  8943  	require.NoError(err)
  8944  	job.Stop = true
  8945  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1200, job)
  8946  	require.NoError(err)
  8947  
  8948  	// Ensure:
  8949  	// * the scaling policy was deleted
  8950  	// * the watches were fired
  8951  	// * the table index was advanced
  8952  	require.True(watchFired(wsGet))
  8953  	require.True(watchFired(wsList))
  8954  	out, err = state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
  8955  	require.NoError(err)
  8956  	require.Nil(out)
  8957  	index, err := state.Index("scaling_policy")
  8958  	require.NoError(err)
  8959  	require.GreaterOrEqual(index, uint64(1200))
  8960  }
  8961  
  8962  func TestStateStore_UnstopJob_UpsertScalingPolicies(t *testing.T) {
  8963  	t.Parallel()
  8964  
  8965  	require := require.New(t)
  8966  
  8967  	state := testStateStore(t)
  8968  
  8969  	job, policy := mock.JobWithScalingPolicy()
  8970  	job.Stop = true
  8971  
  8972  	// establish watcher, verify there are no scaling policies yet
  8973  	ws := memdb.NewWatchSet()
  8974  	list, err := state.ScalingPolicies(ws)
  8975  	require.NoError(err)
  8976  	require.Nil(list.Next())
  8977  
  8978  	// upsert a stopped job, verify that we don't fire the watcher or add any scaling policies
  8979  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  8980  	require.NoError(err)
  8981  	require.True(watchFired(ws))
  8982  	list, err = state.ScalingPolicies(ws)
  8983  	require.NoError(err)
  8984  	require.NotNil(list.Next())
  8985  
  8986  	// Establish a new watchset
  8987  	ws = memdb.NewWatchSet()
  8988  	_, err = state.ScalingPolicies(ws)
  8989  	require.NoError(err)
  8990  	// Unstop this job, say you'll run it again...
  8991  	job.Stop = false
  8992  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1100, job)
  8993  	require.NoError(err)
  8994  
  8995  	// Ensure the scaling policy still exists, watch was not fired, index was not advanced
  8996  	out, err := state.ScalingPolicyByTargetAndType(nil, policy.Target, policy.Type)
  8997  	require.NoError(err)
  8998  	require.NotNil(out)
  8999  	index, err := state.Index("scaling_policy")
  9000  	require.NoError(err)
  9001  	require.EqualValues(index, 1000)
  9002  	require.False(watchFired(ws))
  9003  }
  9004  
  9005  func TestStateStore_DeleteJob_DeleteScalingPolicies(t *testing.T) {
  9006  	t.Parallel()
  9007  
  9008  	require := require.New(t)
  9009  
  9010  	state := testStateStore(t)
  9011  
  9012  	job := mock.Job()
  9013  
  9014  	err := state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  9015  	require.NoError(err)
  9016  
  9017  	policy := mock.ScalingPolicy()
  9018  	policy.Target[structs.ScalingTargetJob] = job.ID
  9019  	err = state.UpsertScalingPolicies(1001, []*structs.ScalingPolicy{policy})
  9020  	require.NoError(err)
  9021  
  9022  	// Delete the job
  9023  	err = state.DeleteJob(1002, job.Namespace, job.ID)
  9024  	require.NoError(err)
  9025  
  9026  	// Ensure the scaling policy was deleted
  9027  	ws := memdb.NewWatchSet()
  9028  	out, err := state.ScalingPolicyByTargetAndType(ws, policy.Target, policy.Type)
  9029  	require.NoError(err)
  9030  	require.Nil(out)
  9031  	index, err := state.Index("scaling_policy")
  9032  	require.NoError(err)
  9033  	require.True(index > 1001)
  9034  }
  9035  
  9036  func TestStateStore_DeleteJob_DeleteScalingPoliciesPrefixBug(t *testing.T) {
  9037  	t.Parallel()
  9038  
  9039  	require := require.New(t)
  9040  
  9041  	state := testStateStore(t)
  9042  
  9043  	job := mock.Job()
  9044  	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 1000, job))
  9045  	job2 := job.Copy()
  9046  	job2.ID = job.ID + "-but-longer"
  9047  	require.NoError(state.UpsertJob(structs.MsgTypeTestSetup, 1001, job2))
  9048  
  9049  	policy := mock.ScalingPolicy()
  9050  	policy.Target[structs.ScalingTargetJob] = job.ID
  9051  	policy2 := mock.ScalingPolicy()
  9052  	policy2.Target[structs.ScalingTargetJob] = job2.ID
  9053  	require.NoError(state.UpsertScalingPolicies(1002, []*structs.ScalingPolicy{policy, policy2}))
  9054  
  9055  	// Delete job with the shorter prefix-ID
  9056  	require.NoError(state.DeleteJob(1003, job.Namespace, job.ID))
  9057  
  9058  	// Ensure only the associated scaling policy was deleted, not the one matching the job with the longer ID
  9059  	out, err := state.ScalingPolicyByID(nil, policy.ID)
  9060  	require.NoError(err)
  9061  	require.Nil(out)
  9062  	out, err = state.ScalingPolicyByID(nil, policy2.ID)
  9063  	require.NoError(err)
  9064  	require.NotNil(out)
  9065  }
  9066  
  9067  // This test ensures that deleting a job that doesn't have any scaling policies
  9068  // will not cause the scaling_policy table index to increase, on either job
  9069  // registration or deletion.
  9070  func TestStateStore_DeleteJob_ScalingPolicyIndexNoop(t *testing.T) {
  9071  	t.Parallel()
  9072  
  9073  	require := require.New(t)
  9074  
  9075  	state := testStateStore(t)
  9076  
  9077  	job := mock.Job()
  9078  
  9079  	prevIndex, err := state.Index("scaling_policy")
  9080  	require.NoError(err)
  9081  
  9082  	err = state.UpsertJob(structs.MsgTypeTestSetup, 1000, job)
  9083  	require.NoError(err)
  9084  
  9085  	newIndex, err := state.Index("scaling_policy")
  9086  	require.NoError(err)
  9087  	require.Equal(prevIndex, newIndex)
  9088  
  9089  	// Delete the job
  9090  	err = state.DeleteJob(1002, job.Namespace, job.ID)
  9091  	require.NoError(err)
  9092  
  9093  	newIndex, err = state.Index("scaling_policy")
  9094  	require.NoError(err)
  9095  	require.Equal(prevIndex, newIndex)
  9096  }
  9097  
  9098  func TestStateStore_ScalingPoliciesByType(t *testing.T) {
  9099  	t.Parallel()
  9100  
  9101  	require := require.New(t)
  9102  
  9103  	state := testStateStore(t)
  9104  
  9105  	// Create scaling policies of different types
  9106  	pHorzA := mock.ScalingPolicy()
  9107  	pHorzA.Type = structs.ScalingPolicyTypeHorizontal
  9108  	pHorzB := mock.ScalingPolicy()
  9109  	pHorzB.Type = structs.ScalingPolicyTypeHorizontal
  9110  
  9111  	pOther1 := mock.ScalingPolicy()
  9112  	pOther1.Type = "other-type-1"
  9113  
  9114  	pOther2 := mock.ScalingPolicy()
  9115  	pOther2.Type = "other-type-2"
  9116  
  9117  	// Create search routine
  9118  	search := func(t string) (found []string) {
  9119  		found = []string{}
  9120  		iter, err := state.ScalingPoliciesByTypePrefix(nil, t)
  9121  		require.NoError(err)
  9122  
  9123  		for raw := iter.Next(); raw != nil; raw = iter.Next() {
  9124  			found = append(found, raw.(*structs.ScalingPolicy).Type)
  9125  		}
  9126  		return
  9127  	}
  9128  
  9129  	// Create the policies
  9130  	var baseIndex uint64 = 1000
  9131  	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{pHorzA, pHorzB, pOther1, pOther2})
  9132  	require.NoError(err)
  9133  
  9134  	// Check if we can read horizontal policies
  9135  	expect := []string{pHorzA.Type, pHorzB.Type}
  9136  	actual := search(structs.ScalingPolicyTypeHorizontal)
  9137  	require.ElementsMatch(expect, actual)
  9138  
  9139  	// Check if we can read policies of other types
  9140  	expect = []string{pOther1.Type}
  9141  	actual = search("other-type-1")
  9142  	require.ElementsMatch(expect, actual)
  9143  
  9144  	// Check that we can read policies by prefix
  9145  	expect = []string{"other-type-1", "other-type-2"}
  9146  	actual = search("other-type")
  9147  	require.Equal(expect, actual)
  9148  
  9149  	// Check for empty result
  9150  	expect = []string{}
  9151  	actual = search("non-existing")
  9152  	require.ElementsMatch(expect, actual)
  9153  }
  9154  
  9155  func TestStateStore_ScalingPoliciesByTypePrefix(t *testing.T) {
  9156  	t.Parallel()
  9157  
  9158  	require := require.New(t)
  9159  
  9160  	state := testStateStore(t)
  9161  
  9162  	// Create scaling policies of different types
  9163  	pHorzA := mock.ScalingPolicy()
  9164  	pHorzA.Type = structs.ScalingPolicyTypeHorizontal
  9165  	pHorzB := mock.ScalingPolicy()
  9166  	pHorzB.Type = structs.ScalingPolicyTypeHorizontal
  9167  
  9168  	pOther1 := mock.ScalingPolicy()
  9169  	pOther1.Type = "other-type-1"
  9170  
  9171  	pOther2 := mock.ScalingPolicy()
  9172  	pOther2.Type = "other-type-2"
  9173  
  9174  	// Create search routine
  9175  	search := func(t string) (count int, found []string, err error) {
  9176  		found = []string{}
  9177  		iter, err := state.ScalingPoliciesByTypePrefix(nil, t)
  9178  		if err != nil {
  9179  			return
  9180  		}
  9181  
  9182  		for raw := iter.Next(); raw != nil; raw = iter.Next() {
  9183  			count++
  9184  			found = append(found, raw.(*structs.ScalingPolicy).Type)
  9185  		}
  9186  		return
  9187  	}
  9188  
  9189  	// Create the policies
  9190  	var baseIndex uint64 = 1000
  9191  	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{pHorzA, pHorzB, pOther1, pOther2})
  9192  	require.NoError(err)
  9193  
  9194  	// Check if we can read horizontal policies
  9195  	expect := []string{pHorzA.Type, pHorzB.Type}
  9196  	count, found, err := search("h")
  9197  
  9198  	sort.Strings(found)
  9199  	sort.Strings(expect)
  9200  
  9201  	require.NoError(err)
  9202  	require.Equal(expect, found)
  9203  	require.Equal(2, count)
  9204  
  9205  	// Check if we can read other prefix policies
  9206  	expect = []string{pOther1.Type, pOther2.Type}
  9207  	count, found, err = search("other")
  9208  
  9209  	sort.Strings(found)
  9210  	sort.Strings(expect)
  9211  
  9212  	require.NoError(err)
  9213  	require.Equal(expect, found)
  9214  	require.Equal(2, count)
  9215  
  9216  	// Check for empty result
  9217  	expect = []string{}
  9218  	count, found, err = search("non-existing")
  9219  
  9220  	sort.Strings(found)
  9221  	sort.Strings(expect)
  9222  
  9223  	require.NoError(err)
  9224  	require.Equal(expect, found)
  9225  	require.Equal(0, count)
  9226  }
  9227  
  9228  func TestStateStore_ScalingPoliciesByJob(t *testing.T) {
  9229  	t.Parallel()
  9230  
  9231  	require := require.New(t)
  9232  
  9233  	state := testStateStore(t)
  9234  	policyA := mock.ScalingPolicy()
  9235  	policyB1 := mock.ScalingPolicy()
  9236  	policyB2 := mock.ScalingPolicy()
  9237  	policyB1.Target[structs.ScalingTargetJob] = policyB2.Target[structs.ScalingTargetJob]
  9238  
  9239  	// Create the policies
  9240  	var baseIndex uint64 = 1000
  9241  	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{policyA, policyB1, policyB2})
  9242  	require.NoError(err)
  9243  
  9244  	iter, err := state.ScalingPoliciesByJob(nil,
  9245  		policyA.Target[structs.ScalingTargetNamespace],
  9246  		policyA.Target[structs.ScalingTargetJob], "")
  9247  	require.NoError(err)
  9248  
  9249  	// Ensure we see expected policies
  9250  	count := 0
  9251  	found := []string{}
  9252  	for {
  9253  		raw := iter.Next()
  9254  		if raw == nil {
  9255  			break
  9256  		}
  9257  		count++
  9258  		found = append(found, raw.(*structs.ScalingPolicy).Target[structs.ScalingTargetGroup])
  9259  	}
  9260  	require.Equal(1, count)
  9261  	sort.Strings(found)
  9262  	expect := []string{policyA.Target[structs.ScalingTargetGroup]}
  9263  	sort.Strings(expect)
  9264  	require.Equal(expect, found)
  9265  
  9266  	iter, err = state.ScalingPoliciesByJob(nil,
  9267  		policyB1.Target[structs.ScalingTargetNamespace],
  9268  		policyB1.Target[structs.ScalingTargetJob], "")
  9269  	require.NoError(err)
  9270  
  9271  	// Ensure we see expected policies
  9272  	count = 0
  9273  	found = []string{}
  9274  	for {
  9275  		raw := iter.Next()
  9276  		if raw == nil {
  9277  			break
  9278  		}
  9279  		count++
  9280  		found = append(found, raw.(*structs.ScalingPolicy).Target[structs.ScalingTargetGroup])
  9281  	}
  9282  	require.Equal(2, count)
  9283  	sort.Strings(found)
  9284  	expect = []string{
  9285  		policyB1.Target[structs.ScalingTargetGroup],
  9286  		policyB2.Target[structs.ScalingTargetGroup],
  9287  	}
  9288  	sort.Strings(expect)
  9289  	require.Equal(expect, found)
  9290  }
  9291  
  9292  func TestStateStore_ScalingPoliciesByJob_PrefixBug(t *testing.T) {
  9293  	t.Parallel()
  9294  
  9295  	require := require.New(t)
  9296  
  9297  	jobPrefix := "job-name-" + uuid.Generate()
  9298  
  9299  	state := testStateStore(t)
  9300  	policy1 := mock.ScalingPolicy()
  9301  	policy1.Target[structs.ScalingTargetJob] = jobPrefix
  9302  	policy2 := mock.ScalingPolicy()
  9303  	policy2.Target[structs.ScalingTargetJob] = jobPrefix + "-more"
  9304  
  9305  	// Create the policies
  9306  	var baseIndex uint64 = 1000
  9307  	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{policy1, policy2})
  9308  	require.NoError(err)
  9309  
  9310  	iter, err := state.ScalingPoliciesByJob(nil,
  9311  		policy1.Target[structs.ScalingTargetNamespace],
  9312  		jobPrefix, "")
  9313  	require.NoError(err)
  9314  
  9315  	// Ensure we see expected policies
  9316  	count := 0
  9317  	found := []string{}
  9318  	for {
  9319  		raw := iter.Next()
  9320  		if raw == nil {
  9321  			break
  9322  		}
  9323  		count++
  9324  		found = append(found, raw.(*structs.ScalingPolicy).ID)
  9325  	}
  9326  	require.Equal(1, count)
  9327  	expect := []string{policy1.ID}
  9328  	require.Equal(expect, found)
  9329  }
  9330  
  9331  func TestStateStore_ScalingPolicyByTargetAndType(t *testing.T) {
  9332  	t.Parallel()
  9333  
  9334  	require := require.New(t)
  9335  
  9336  	state := testStateStore(t)
  9337  
  9338  	// Create scaling policies
  9339  	policyA := mock.ScalingPolicy()
  9340  	// Same target, different type
  9341  	policyB := mock.ScalingPolicy()
  9342  	policyC := mock.ScalingPolicy()
  9343  	for k, v := range policyB.Target {
  9344  		policyC.Target[k] = v
  9345  	}
  9346  	policyC.Type = "other-type"
  9347  
  9348  	// Create the policies
  9349  	var baseIndex uint64 = 1000
  9350  	err := state.UpsertScalingPolicies(baseIndex, []*structs.ScalingPolicy{policyA, policyB, policyC})
  9351  	require.NoError(err)
  9352  
  9353  	// Check if we can retrieve the right policies
  9354  	found, err := state.ScalingPolicyByTargetAndType(nil, policyA.Target, policyA.Type)
  9355  	require.NoError(err)
  9356  	require.Equal(policyA, found)
  9357  
  9358  	// Check for wrong type
  9359  	found, err = state.ScalingPolicyByTargetAndType(nil, policyA.Target, "wrong_type")
  9360  	require.NoError(err)
  9361  	require.Nil(found)
  9362  
  9363  	// Check for same target but different type
  9364  	found, err = state.ScalingPolicyByTargetAndType(nil, policyB.Target, policyB.Type)
  9365  	require.NoError(err)
  9366  	require.Equal(policyB, found)
  9367  
  9368  	found, err = state.ScalingPolicyByTargetAndType(nil, policyB.Target, policyC.Type)
  9369  	require.NoError(err)
  9370  	require.Equal(policyC, found)
  9371  }
  9372  
  9373  func TestStateStore_UpsertScalingEvent(t *testing.T) {
  9374  	t.Parallel()
  9375  	require := require.New(t)
  9376  
  9377  	state := testStateStore(t)
  9378  	job := mock.Job()
  9379  	groupName := job.TaskGroups[0].Name
  9380  
  9381  	newEvent := structs.NewScalingEvent("message 1").SetMeta(map[string]interface{}{
  9382  		"a": 1,
  9383  	})
  9384  
  9385  	wsAll := memdb.NewWatchSet()
  9386  	all, err := state.ScalingEvents(wsAll)
  9387  	require.NoError(err)
  9388  	require.Nil(all.Next())
  9389  
  9390  	ws := memdb.NewWatchSet()
  9391  	out, _, err := state.ScalingEventsByJob(ws, job.Namespace, job.ID)
  9392  	require.NoError(err)
  9393  	require.Nil(out)
  9394  
  9395  	err = state.UpsertScalingEvent(1000, &structs.ScalingEventRequest{
  9396  		Namespace:    job.Namespace,
  9397  		JobID:        job.ID,
  9398  		TaskGroup:    groupName,
  9399  		ScalingEvent: newEvent,
  9400  	})
  9401  	require.NoError(err)
  9402  	require.True(watchFired(ws))
  9403  	require.True(watchFired(wsAll))
  9404  
  9405  	ws = memdb.NewWatchSet()
  9406  	out, eventsIndex, err := state.ScalingEventsByJob(ws, job.Namespace, job.ID)
  9407  	require.NoError(err)
  9408  	require.Equal(map[string][]*structs.ScalingEvent{
  9409  		groupName: {newEvent},
  9410  	}, out)
  9411  	require.EqualValues(eventsIndex, 1000)
  9412  
  9413  	iter, err := state.ScalingEvents(ws)
  9414  	require.NoError(err)
  9415  
  9416  	count := 0
  9417  	jobsReturned := []string{}
  9418  	var jobEvents *structs.JobScalingEvents
  9419  	for {
  9420  		raw := iter.Next()
  9421  		if raw == nil {
  9422  			break
  9423  		}
  9424  		jobEvents = raw.(*structs.JobScalingEvents)
  9425  		jobsReturned = append(jobsReturned, jobEvents.JobID)
  9426  		count++
  9427  	}
  9428  	require.Equal(1, count)
  9429  	require.EqualValues(jobEvents.ModifyIndex, 1000)
  9430  	require.EqualValues(jobEvents.ScalingEvents[groupName][0].CreateIndex, 1000)
  9431  
  9432  	index, err := state.Index("scaling_event")
  9433  	require.NoError(err)
  9434  	require.ElementsMatch([]string{job.ID}, jobsReturned)
  9435  	require.Equal(map[string][]*structs.ScalingEvent{
  9436  		groupName: {newEvent},
  9437  	}, jobEvents.ScalingEvents)
  9438  	require.EqualValues(1000, index)
  9439  	require.False(watchFired(ws))
  9440  }
  9441  
  9442  func TestStateStore_UpsertScalingEvent_LimitAndOrder(t *testing.T) {
  9443  	t.Parallel()
  9444  	require := require.New(t)
  9445  
  9446  	state := testStateStore(t)
  9447  	namespace := uuid.Generate()
  9448  	jobID := uuid.Generate()
  9449  	group1 := uuid.Generate()
  9450  	group2 := uuid.Generate()
  9451  
  9452  	index := uint64(1000)
  9453  	for i := 1; i <= structs.JobTrackedScalingEvents+10; i++ {
  9454  		newEvent := structs.NewScalingEvent("").SetMeta(map[string]interface{}{
  9455  			"i":     i,
  9456  			"group": group1,
  9457  		})
  9458  		err := state.UpsertScalingEvent(index, &structs.ScalingEventRequest{
  9459  			Namespace:    namespace,
  9460  			JobID:        jobID,
  9461  			TaskGroup:    group1,
  9462  			ScalingEvent: newEvent,
  9463  		})
  9464  		index++
  9465  		require.NoError(err)
  9466  
  9467  		newEvent = structs.NewScalingEvent("").SetMeta(map[string]interface{}{
  9468  			"i":     i,
  9469  			"group": group2,
  9470  		})
  9471  		err = state.UpsertScalingEvent(index, &structs.ScalingEventRequest{
  9472  			Namespace:    namespace,
  9473  			JobID:        jobID,
  9474  			TaskGroup:    group2,
  9475  			ScalingEvent: newEvent,
  9476  		})
  9477  		index++
  9478  		require.NoError(err)
  9479  	}
  9480  
  9481  	out, _, err := state.ScalingEventsByJob(nil, namespace, jobID)
  9482  	require.NoError(err)
  9483  	require.Len(out, 2)
  9484  
  9485  	expectedEvents := []int{}
  9486  	for i := structs.JobTrackedScalingEvents; i > 0; i-- {
  9487  		expectedEvents = append(expectedEvents, i+10)
  9488  	}
  9489  
  9490  	// checking order and content
  9491  	require.Len(out[group1], structs.JobTrackedScalingEvents)
  9492  	actualEvents := []int{}
  9493  	for _, event := range out[group1] {
  9494  		require.Equal(group1, event.Meta["group"])
  9495  		actualEvents = append(actualEvents, event.Meta["i"].(int))
  9496  	}
  9497  	require.Equal(expectedEvents, actualEvents)
  9498  
  9499  	// checking order and content
  9500  	require.Len(out[group2], structs.JobTrackedScalingEvents)
  9501  	actualEvents = []int{}
  9502  	for _, event := range out[group2] {
  9503  		require.Equal(group2, event.Meta["group"])
  9504  		actualEvents = append(actualEvents, event.Meta["i"].(int))
  9505  	}
  9506  	require.Equal(expectedEvents, actualEvents)
  9507  }
  9508  
  9509  func TestStateStore_RestoreScalingEvents(t *testing.T) {
  9510  	t.Parallel()
  9511  	require := require.New(t)
  9512  
  9513  	state := testStateStore(t)
  9514  	jobScalingEvents := &structs.JobScalingEvents{
  9515  		Namespace: uuid.Generate(),
  9516  		JobID:     uuid.Generate(),
  9517  		ScalingEvents: map[string][]*structs.ScalingEvent{
  9518  			uuid.Generate(): {
  9519  				structs.NewScalingEvent(uuid.Generate()),
  9520  			},
  9521  		},
  9522  	}
  9523  
  9524  	restore, err := state.Restore()
  9525  	require.NoError(err)
  9526  
  9527  	err = restore.ScalingEventsRestore(jobScalingEvents)
  9528  	require.NoError(err)
  9529  	require.NoError(restore.Commit())
  9530  
  9531  	ws := memdb.NewWatchSet()
  9532  	out, _, err := state.ScalingEventsByJob(ws, jobScalingEvents.Namespace,
  9533  		jobScalingEvents.JobID)
  9534  	require.NoError(err)
  9535  	require.NotNil(out)
  9536  	require.EqualValues(jobScalingEvents.ScalingEvents, out)
  9537  }
  9538  
  9539  func TestStateStore_Abandon(t *testing.T) {
  9540  	t.Parallel()
  9541  
  9542  	s := testStateStore(t)
  9543  	abandonCh := s.AbandonCh()
  9544  	s.Abandon()
  9545  	select {
  9546  	case <-abandonCh:
  9547  	default:
  9548  		t.Fatalf("bad")
  9549  	}
  9550  }
  9551  
  9552  // Verifies that an error is returned when an allocation doesn't exist in the state store.
  9553  func TestStateSnapshot_DenormalizeAllocationDiffSlice_AllocDoesNotExist(t *testing.T) {
  9554  	t.Parallel()
  9555  
  9556  	state := testStateStore(t)
  9557  	alloc := mock.Alloc()
  9558  	require := require.New(t)
  9559  
  9560  	// Insert job
  9561  	err := state.UpsertJob(structs.MsgTypeTestSetup, 999, alloc.Job)
  9562  	require.NoError(err)
  9563  
  9564  	allocDiffs := []*structs.AllocationDiff{
  9565  		{
  9566  			ID: alloc.ID,
  9567  		},
  9568  	}
  9569  
  9570  	snap, err := state.Snapshot()
  9571  	require.NoError(err)
  9572  
  9573  	denormalizedAllocs, err := snap.DenormalizeAllocationDiffSlice(allocDiffs)
  9574  
  9575  	require.EqualError(err, fmt.Sprintf("alloc %v doesn't exist", alloc.ID))
  9576  	require.Nil(denormalizedAllocs)
  9577  }
  9578  
  9579  // TestStateStore_SnapshotMinIndex_OK asserts StateStore.SnapshotMinIndex blocks
  9580  // until the StateStore's latest index is >= the requested index.
  9581  func TestStateStore_SnapshotMinIndex_OK(t *testing.T) {
  9582  	t.Parallel()
  9583  
  9584  	s := testStateStore(t)
  9585  	index, err := s.LatestIndex()
  9586  	require.NoError(t, err)
  9587  
  9588  	node := mock.Node()
  9589  	require.NoError(t, s.UpsertNode(structs.MsgTypeTestSetup, index+1, node))
  9590  
  9591  	// Assert SnapshotMinIndex returns immediately if index < latest index
  9592  	ctx, cancel := context.WithTimeout(context.Background(), 0)
  9593  	snap, err := s.SnapshotMinIndex(ctx, index)
  9594  	cancel()
  9595  	require.NoError(t, err)
  9596  
  9597  	snapIndex, err := snap.LatestIndex()
  9598  	require.NoError(t, err)
  9599  	if snapIndex <= index {
  9600  		require.Fail(t, "snapshot index should be greater than index")
  9601  	}
  9602  
  9603  	// Assert SnapshotMinIndex returns immediately if index == latest index
  9604  	ctx, cancel = context.WithTimeout(context.Background(), 0)
  9605  	snap, err = s.SnapshotMinIndex(ctx, index+1)
  9606  	cancel()
  9607  	require.NoError(t, err)
  9608  
  9609  	snapIndex, err = snap.LatestIndex()
  9610  	require.NoError(t, err)
  9611  	require.Equal(t, snapIndex, index+1)
  9612  
  9613  	// Assert SnapshotMinIndex blocks if index > latest index
  9614  	errCh := make(chan error, 1)
  9615  	ctx, cancel = context.WithTimeout(context.Background(), 10*time.Second)
  9616  	defer cancel()
  9617  	go func() {
  9618  		defer close(errCh)
  9619  		waitIndex := index + 2
  9620  		snap, err := s.SnapshotMinIndex(ctx, waitIndex)
  9621  		if err != nil {
  9622  			errCh <- err
  9623  			return
  9624  		}
  9625  
  9626  		snapIndex, err := snap.LatestIndex()
  9627  		if err != nil {
  9628  			errCh <- err
  9629  			return
  9630  		}
  9631  
  9632  		if snapIndex < waitIndex {
  9633  			errCh <- fmt.Errorf("snapshot index < wait index: %d < %d", snapIndex, waitIndex)
  9634  			return
  9635  		}
  9636  	}()
  9637  
  9638  	select {
  9639  	case err := <-errCh:
  9640  		require.NoError(t, err)
  9641  	case <-time.After(500 * time.Millisecond):
  9642  		// Let it block for a bit before unblocking by upserting
  9643  	}
  9644  
  9645  	node.Name = "hal"
  9646  	require.NoError(t, s.UpsertNode(structs.MsgTypeTestSetup, index+2, node))
  9647  
  9648  	select {
  9649  	case err := <-errCh:
  9650  		require.NoError(t, err)
  9651  	case <-time.After(5 * time.Second):
  9652  		require.Fail(t, "timed out waiting for SnapshotMinIndex to unblock")
  9653  	}
  9654  }
  9655  
  9656  // TestStateStore_SnapshotMinIndex_Timeout asserts StateStore.SnapshotMinIndex
  9657  // returns an error if the desired index is not reached within the deadline.
  9658  func TestStateStore_SnapshotMinIndex_Timeout(t *testing.T) {
  9659  	t.Parallel()
  9660  
  9661  	s := testStateStore(t)
  9662  	index, err := s.LatestIndex()
  9663  	require.NoError(t, err)
  9664  
  9665  	// Assert SnapshotMinIndex blocks if index > latest index
  9666  	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
  9667  	defer cancel()
  9668  	snap, err := s.SnapshotMinIndex(ctx, index+1)
  9669  	require.EqualError(t, err, context.DeadlineExceeded.Error())
  9670  	require.Nil(t, snap)
  9671  }
  9672  
  9673  // watchFired is a helper for unit tests that returns if the given watch set
  9674  // fired (it doesn't care which watch actually fired). This uses a fixed
  9675  // timeout since we already expect the event happened before calling this and
  9676  // just need to distinguish a fire from a timeout. We do need a little time to
  9677  // allow the watch to set up any goroutines, though.
  9678  func watchFired(ws memdb.WatchSet) bool {
  9679  	timedOut := ws.Watch(time.After(50 * time.Millisecond))
  9680  	return !timedOut
  9681  }
  9682  
  9683  // NodeIDSort is used to sort nodes by ID
  9684  type NodeIDSort []*structs.Node
  9685  
  9686  func (n NodeIDSort) Len() int {
  9687  	return len(n)
  9688  }
  9689  
  9690  func (n NodeIDSort) Less(i, j int) bool {
  9691  	return n[i].ID < n[j].ID
  9692  }
  9693  
  9694  func (n NodeIDSort) Swap(i, j int) {
  9695  	n[i], n[j] = n[j], n[i]
  9696  }
  9697  
  9698  // JobIDis used to sort jobs by id
  9699  type JobIDSort []*structs.Job
  9700  
  9701  func (n JobIDSort) Len() int {
  9702  	return len(n)
  9703  }
  9704  
  9705  func (n JobIDSort) Less(i, j int) bool {
  9706  	return n[i].ID < n[j].ID
  9707  }
  9708  
  9709  func (n JobIDSort) Swap(i, j int) {
  9710  	n[i], n[j] = n[j], n[i]
  9711  }
  9712  
  9713  // EvalIDis used to sort evals by id
  9714  type EvalIDSort []*structs.Evaluation
  9715  
  9716  func (n EvalIDSort) Len() int {
  9717  	return len(n)
  9718  }
  9719  
  9720  func (n EvalIDSort) Less(i, j int) bool {
  9721  	return n[i].ID < n[j].ID
  9722  }
  9723  
  9724  func (n EvalIDSort) Swap(i, j int) {
  9725  	n[i], n[j] = n[j], n[i]
  9726  }
  9727  
  9728  // AllocIDsort used to sort allocations by id
  9729  type AllocIDSort []*structs.Allocation
  9730  
  9731  func (n AllocIDSort) Len() int {
  9732  	return len(n)
  9733  }
  9734  
  9735  func (n AllocIDSort) Less(i, j int) bool {
  9736  	return n[i].ID < n[j].ID
  9737  }
  9738  
  9739  func (n AllocIDSort) Swap(i, j int) {
  9740  	n[i], n[j] = n[j], n[i]
  9741  }