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