github.com/zoomfoo/nomad@v0.8.5-0.20180907175415-f28fd3a1a056/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  	memdb "github.com/hashicorp/go-memdb"
    13  	"github.com/hashicorp/nomad/helper"
    14  	"github.com/hashicorp/nomad/helper/uuid"
    15  	"github.com/hashicorp/nomad/nomad/mock"
    16  	"github.com/hashicorp/nomad/nomad/structs"
    17  	"github.com/stretchr/testify/assert"
    18  	"github.com/stretchr/testify/require"
    19  )
    20  
    21  func testStateStore(t *testing.T) *StateStore {
    22  	return TestStateStore(t)
    23  }
    24  
    25  func TestStateStore_Blocking_Error(t *testing.T) {
    26  	t.Parallel()
    27  	expected := fmt.Errorf("test error")
    28  	errFn := func(memdb.WatchSet, *StateStore) (interface{}, uint64, error) {
    29  		return nil, 0, expected
    30  	}
    31  
    32  	state := testStateStore(t)
    33  	_, idx, err := state.BlockingQuery(errFn, 10, context.Background())
    34  	assert.EqualError(t, err, expected.Error())
    35  	assert.Zero(t, idx)
    36  }
    37  
    38  func TestStateStore_Blocking_Timeout(t *testing.T) {
    39  	t.Parallel()
    40  	noopFn := func(memdb.WatchSet, *StateStore) (interface{}, uint64, error) {
    41  		return nil, 5, nil
    42  	}
    43  
    44  	state := testStateStore(t)
    45  	timeout := time.Now().Add(250 * time.Millisecond)
    46  	deadlineCtx, cancel := context.WithDeadline(context.Background(), timeout)
    47  	defer cancel()
    48  
    49  	_, idx, err := state.BlockingQuery(noopFn, 10, deadlineCtx)
    50  	assert.EqualError(t, err, context.DeadlineExceeded.Error())
    51  	assert.EqualValues(t, 5, idx)
    52  	assert.WithinDuration(t, timeout, time.Now(), 100*time.Millisecond)
    53  }
    54  
    55  func TestStateStore_Blocking_MinQuery(t *testing.T) {
    56  	node := mock.Node()
    57  	count := 0
    58  	queryFn := func(ws memdb.WatchSet, s *StateStore) (interface{}, uint64, error) {
    59  		_, err := s.NodeByID(ws, node.ID)
    60  		if err != nil {
    61  			return nil, 0, err
    62  		}
    63  
    64  		count++
    65  		if count == 1 {
    66  			return false, 5, nil
    67  		} else if count > 2 {
    68  			return false, 20, fmt.Errorf("called too many times")
    69  		}
    70  
    71  		return true, 11, nil
    72  	}
    73  
    74  	state := testStateStore(t)
    75  	timeout := time.Now().Add(100 * time.Millisecond)
    76  	deadlineCtx, cancel := context.WithDeadline(context.Background(), timeout)
    77  	defer cancel()
    78  
    79  	time.AfterFunc(5*time.Millisecond, func() {
    80  		state.UpsertNode(11, node)
    81  	})
    82  
    83  	resp, idx, err := state.BlockingQuery(queryFn, 10, deadlineCtx)
    84  	if assert.Nil(t, err) {
    85  		assert.Equal(t, 2, count)
    86  		assert.EqualValues(t, 11, idx)
    87  		assert.True(t, resp.(bool))
    88  	}
    89  }
    90  
    91  // This test checks that:
    92  // 1) The job is denormalized
    93  // 2) Allocations are created
    94  func TestStateStore_UpsertPlanResults_AllocationsCreated_Denormalized(t *testing.T) {
    95  	state := testStateStore(t)
    96  	alloc := mock.Alloc()
    97  	job := alloc.Job
    98  	alloc.Job = nil
    99  
   100  	if err := state.UpsertJob(999, job); err != nil {
   101  		t.Fatalf("err: %v", err)
   102  	}
   103  
   104  	eval := mock.Eval()
   105  	eval.JobID = job.ID
   106  
   107  	// Create an eval
   108  	if err := state.UpsertEvals(1, []*structs.Evaluation{eval}); err != nil {
   109  		t.Fatalf("err: %v", err)
   110  	}
   111  
   112  	// Create a plan result
   113  	res := structs.ApplyPlanResultsRequest{
   114  		AllocUpdateRequest: structs.AllocUpdateRequest{
   115  			Alloc: []*structs.Allocation{alloc},
   116  			Job:   job,
   117  		},
   118  		EvalID: eval.ID,
   119  	}
   120  	assert := assert.New(t)
   121  	err := state.UpsertPlanResults(1000, &res)
   122  	assert.Nil(err)
   123  
   124  	ws := memdb.NewWatchSet()
   125  	out, err := state.AllocByID(ws, alloc.ID)
   126  	assert.Nil(err)
   127  	assert.Equal(alloc, out)
   128  
   129  	index, err := state.Index("allocs")
   130  	assert.Nil(err)
   131  	assert.EqualValues(1000, index)
   132  
   133  	if watchFired(ws) {
   134  		t.Fatalf("bad")
   135  	}
   136  
   137  	evalOut, err := state.EvalByID(ws, eval.ID)
   138  	assert.Nil(err)
   139  	assert.NotNil(evalOut)
   140  	assert.EqualValues(1000, evalOut.ModifyIndex)
   141  }
   142  
   143  // This test checks that the deployment is created and allocations count towards
   144  // the deployment
   145  func TestStateStore_UpsertPlanResults_Deployment(t *testing.T) {
   146  	state := testStateStore(t)
   147  	alloc := mock.Alloc()
   148  	alloc2 := mock.Alloc()
   149  	job := alloc.Job
   150  	alloc.Job = nil
   151  	alloc2.Job = nil
   152  
   153  	d := mock.Deployment()
   154  	alloc.DeploymentID = d.ID
   155  	alloc2.DeploymentID = d.ID
   156  
   157  	if err := state.UpsertJob(999, job); err != nil {
   158  		t.Fatalf("err: %v", err)
   159  	}
   160  
   161  	eval := mock.Eval()
   162  	eval.JobID = job.ID
   163  
   164  	// Create an eval
   165  	if err := state.UpsertEvals(1, []*structs.Evaluation{eval}); err != nil {
   166  		t.Fatalf("err: %v", err)
   167  	}
   168  
   169  	// Create a plan result
   170  	res := structs.ApplyPlanResultsRequest{
   171  		AllocUpdateRequest: structs.AllocUpdateRequest{
   172  			Alloc: []*structs.Allocation{alloc, alloc2},
   173  			Job:   job,
   174  		},
   175  		Deployment: d,
   176  		EvalID:     eval.ID,
   177  	}
   178  
   179  	err := state.UpsertPlanResults(1000, &res)
   180  	if err != nil {
   181  		t.Fatalf("err: %v", err)
   182  	}
   183  
   184  	ws := memdb.NewWatchSet()
   185  	assert := assert.New(t)
   186  	out, err := state.AllocByID(ws, alloc.ID)
   187  	assert.Nil(err)
   188  	assert.Equal(alloc, out)
   189  
   190  	dout, err := state.DeploymentByID(ws, d.ID)
   191  	assert.Nil(err)
   192  	assert.NotNil(dout)
   193  
   194  	tg, ok := dout.TaskGroups[alloc.TaskGroup]
   195  	assert.True(ok)
   196  	assert.NotNil(tg)
   197  	assert.Equal(2, tg.PlacedAllocs)
   198  
   199  	evalOut, err := state.EvalByID(ws, eval.ID)
   200  	assert.Nil(err)
   201  	assert.NotNil(evalOut)
   202  	assert.EqualValues(1000, evalOut.ModifyIndex)
   203  
   204  	if watchFired(ws) {
   205  		t.Fatalf("bad")
   206  	}
   207  
   208  	// Update the allocs to be part of a new deployment
   209  	d2 := d.Copy()
   210  	d2.ID = uuid.Generate()
   211  
   212  	allocNew := alloc.Copy()
   213  	allocNew.DeploymentID = d2.ID
   214  	allocNew2 := alloc2.Copy()
   215  	allocNew2.DeploymentID = d2.ID
   216  
   217  	// Create another plan
   218  	res = structs.ApplyPlanResultsRequest{
   219  		AllocUpdateRequest: structs.AllocUpdateRequest{
   220  			Alloc: []*structs.Allocation{allocNew, allocNew2},
   221  			Job:   job,
   222  		},
   223  		Deployment: d2,
   224  		EvalID:     eval.ID,
   225  	}
   226  
   227  	err = state.UpsertPlanResults(1001, &res)
   228  	if err != nil {
   229  		t.Fatalf("err: %v", err)
   230  	}
   231  
   232  	dout, err = state.DeploymentByID(ws, d2.ID)
   233  	assert.Nil(err)
   234  	assert.NotNil(dout)
   235  
   236  	tg, ok = dout.TaskGroups[alloc.TaskGroup]
   237  	assert.True(ok)
   238  	assert.NotNil(tg)
   239  	assert.Equal(2, tg.PlacedAllocs)
   240  
   241  	evalOut, err = state.EvalByID(ws, eval.ID)
   242  	assert.Nil(err)
   243  	assert.NotNil(evalOut)
   244  	assert.EqualValues(1001, evalOut.ModifyIndex)
   245  }
   246  
   247  // This test checks that deployment updates are applied correctly
   248  func TestStateStore_UpsertPlanResults_DeploymentUpdates(t *testing.T) {
   249  	state := testStateStore(t)
   250  
   251  	// Create a job that applies to all
   252  	job := mock.Job()
   253  	if err := state.UpsertJob(998, job); err != nil {
   254  		t.Fatalf("err: %v", err)
   255  	}
   256  
   257  	// Create a deployment that we will update its status
   258  	doutstanding := mock.Deployment()
   259  	doutstanding.JobID = job.ID
   260  
   261  	if err := state.UpsertDeployment(1000, doutstanding); err != nil {
   262  		t.Fatalf("err: %v", err)
   263  	}
   264  
   265  	eval := mock.Eval()
   266  	eval.JobID = job.ID
   267  
   268  	// Create an eval
   269  	if err := state.UpsertEvals(1, []*structs.Evaluation{eval}); err != nil {
   270  		t.Fatalf("err: %v", err)
   271  	}
   272  	alloc := mock.Alloc()
   273  	alloc.Job = nil
   274  
   275  	dnew := mock.Deployment()
   276  	dnew.JobID = job.ID
   277  	alloc.DeploymentID = dnew.ID
   278  
   279  	// Update the old deployment
   280  	update := &structs.DeploymentStatusUpdate{
   281  		DeploymentID:      doutstanding.ID,
   282  		Status:            "foo",
   283  		StatusDescription: "bar",
   284  	}
   285  
   286  	// Create a plan result
   287  	res := structs.ApplyPlanResultsRequest{
   288  		AllocUpdateRequest: structs.AllocUpdateRequest{
   289  			Alloc: []*structs.Allocation{alloc},
   290  			Job:   job,
   291  		},
   292  		Deployment:        dnew,
   293  		DeploymentUpdates: []*structs.DeploymentStatusUpdate{update},
   294  		EvalID:            eval.ID,
   295  	}
   296  
   297  	err := state.UpsertPlanResults(1000, &res)
   298  	if err != nil {
   299  		t.Fatalf("err: %v", err)
   300  	}
   301  	assert := assert.New(t)
   302  	ws := memdb.NewWatchSet()
   303  
   304  	// Check the deployments are correctly updated.
   305  	dout, err := state.DeploymentByID(ws, dnew.ID)
   306  	assert.Nil(err)
   307  	assert.NotNil(dout)
   308  
   309  	tg, ok := dout.TaskGroups[alloc.TaskGroup]
   310  	assert.True(ok)
   311  	assert.NotNil(tg)
   312  	assert.Equal(1, tg.PlacedAllocs)
   313  
   314  	doutstandingout, err := state.DeploymentByID(ws, doutstanding.ID)
   315  	assert.Nil(err)
   316  	assert.NotNil(doutstandingout)
   317  	assert.Equal(update.Status, doutstandingout.Status)
   318  	assert.Equal(update.StatusDescription, doutstandingout.StatusDescription)
   319  	assert.EqualValues(1000, doutstandingout.ModifyIndex)
   320  
   321  	evalOut, err := state.EvalByID(ws, eval.ID)
   322  	assert.Nil(err)
   323  	assert.NotNil(evalOut)
   324  	assert.EqualValues(1000, evalOut.ModifyIndex)
   325  	if watchFired(ws) {
   326  		t.Fatalf("bad")
   327  	}
   328  }
   329  
   330  func TestStateStore_UpsertDeployment(t *testing.T) {
   331  	state := testStateStore(t)
   332  	deployment := mock.Deployment()
   333  
   334  	// Create a watchset so we can test that upsert fires the watch
   335  	ws := memdb.NewWatchSet()
   336  	_, err := state.DeploymentsByJobID(ws, deployment.Namespace, deployment.ID)
   337  	if err != nil {
   338  		t.Fatalf("bad: %v", err)
   339  	}
   340  
   341  	err = state.UpsertDeployment(1000, deployment)
   342  	if err != nil {
   343  		t.Fatalf("err: %v", err)
   344  	}
   345  	if !watchFired(ws) {
   346  		t.Fatalf("bad")
   347  	}
   348  
   349  	ws = memdb.NewWatchSet()
   350  	out, err := state.DeploymentByID(ws, deployment.ID)
   351  	if err != nil {
   352  		t.Fatalf("err: %v", err)
   353  	}
   354  
   355  	if !reflect.DeepEqual(deployment, out) {
   356  		t.Fatalf("bad: %#v %#v", deployment, out)
   357  	}
   358  
   359  	index, err := state.Index("deployment")
   360  	if err != nil {
   361  		t.Fatalf("err: %v", err)
   362  	}
   363  	if index != 1000 {
   364  		t.Fatalf("bad: %d", index)
   365  	}
   366  
   367  	if watchFired(ws) {
   368  		t.Fatalf("bad")
   369  	}
   370  }
   371  
   372  func TestStateStore_DeleteDeployment(t *testing.T) {
   373  	state := testStateStore(t)
   374  	d1 := mock.Deployment()
   375  	d2 := mock.Deployment()
   376  
   377  	err := state.UpsertDeployment(1000, d1)
   378  	if err != nil {
   379  		t.Fatalf("err: %v", err)
   380  	}
   381  	if err := state.UpsertDeployment(1001, d2); err != nil {
   382  		t.Fatalf("err: %v", err)
   383  	}
   384  
   385  	// Create a watchset so we can test that delete fires the watch
   386  	ws := memdb.NewWatchSet()
   387  	if _, err := state.DeploymentByID(ws, d1.ID); err != nil {
   388  		t.Fatalf("bad: %v", err)
   389  	}
   390  
   391  	err = state.DeleteDeployment(1002, []string{d1.ID, d2.ID})
   392  	if err != nil {
   393  		t.Fatalf("err: %v", err)
   394  	}
   395  
   396  	if !watchFired(ws) {
   397  		t.Fatalf("bad")
   398  	}
   399  
   400  	ws = memdb.NewWatchSet()
   401  	out, err := state.DeploymentByID(ws, d1.ID)
   402  	if err != nil {
   403  		t.Fatalf("err: %v", err)
   404  	}
   405  
   406  	if out != nil {
   407  		t.Fatalf("bad: %#v %#v", d1, out)
   408  	}
   409  
   410  	index, err := state.Index("deployment")
   411  	if err != nil {
   412  		t.Fatalf("err: %v", err)
   413  	}
   414  	if index != 1002 {
   415  		t.Fatalf("bad: %d", index)
   416  	}
   417  
   418  	if watchFired(ws) {
   419  		t.Fatalf("bad")
   420  	}
   421  }
   422  
   423  func TestStateStore_Deployments(t *testing.T) {
   424  	state := testStateStore(t)
   425  	var deployments []*structs.Deployment
   426  
   427  	for i := 0; i < 10; i++ {
   428  		deployment := mock.Deployment()
   429  		deployments = append(deployments, deployment)
   430  
   431  		err := state.UpsertDeployment(1000+uint64(i), deployment)
   432  		if err != nil {
   433  			t.Fatalf("err: %v", err)
   434  		}
   435  	}
   436  
   437  	ws := memdb.NewWatchSet()
   438  	iter, err := state.Deployments(ws)
   439  	if err != nil {
   440  		t.Fatalf("err: %v", err)
   441  	}
   442  
   443  	var out []*structs.Deployment
   444  	for {
   445  		raw := iter.Next()
   446  		if raw == nil {
   447  			break
   448  		}
   449  		out = append(out, raw.(*structs.Deployment))
   450  	}
   451  
   452  	lessThan := func(i, j int) bool {
   453  		return deployments[i].ID < deployments[j].ID
   454  	}
   455  	sort.Slice(deployments, lessThan)
   456  	sort.Slice(out, lessThan)
   457  
   458  	if !reflect.DeepEqual(deployments, out) {
   459  		t.Fatalf("bad: %#v %#v", deployments, out)
   460  	}
   461  
   462  	if watchFired(ws) {
   463  		t.Fatalf("bad")
   464  	}
   465  }
   466  
   467  func TestStateStore_DeploymentsByIDPrefix(t *testing.T) {
   468  	state := testStateStore(t)
   469  	deploy := mock.Deployment()
   470  
   471  	deploy.ID = "11111111-662e-d0ab-d1c9-3e434af7bdb4"
   472  	err := state.UpsertDeployment(1000, deploy)
   473  	if err != nil {
   474  		t.Fatalf("err: %v", err)
   475  	}
   476  
   477  	// Create a watchset so we can test that getters don't cause it to fire
   478  	ws := memdb.NewWatchSet()
   479  	iter, err := state.DeploymentsByIDPrefix(ws, deploy.Namespace, deploy.ID)
   480  	if err != nil {
   481  		t.Fatalf("err: %v", err)
   482  	}
   483  
   484  	gatherDeploys := func(iter memdb.ResultIterator) []*structs.Deployment {
   485  		var deploys []*structs.Deployment
   486  		for {
   487  			raw := iter.Next()
   488  			if raw == nil {
   489  				break
   490  			}
   491  			deploy := raw.(*structs.Deployment)
   492  			deploys = append(deploys, deploy)
   493  		}
   494  		return deploys
   495  	}
   496  
   497  	deploys := gatherDeploys(iter)
   498  	if len(deploys) != 1 {
   499  		t.Fatalf("err: %v", err)
   500  	}
   501  
   502  	if watchFired(ws) {
   503  		t.Fatalf("bad")
   504  	}
   505  
   506  	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "11")
   507  	if err != nil {
   508  		t.Fatalf("err: %v", err)
   509  	}
   510  
   511  	deploys = gatherDeploys(iter)
   512  	if len(deploys) != 1 {
   513  		t.Fatalf("err: %v", err)
   514  	}
   515  
   516  	deploy = mock.Deployment()
   517  	deploy.ID = "11222222-662e-d0ab-d1c9-3e434af7bdb4"
   518  	err = state.UpsertDeployment(1001, deploy)
   519  	if err != nil {
   520  		t.Fatalf("err: %v", err)
   521  	}
   522  
   523  	if !watchFired(ws) {
   524  		t.Fatalf("bad")
   525  	}
   526  
   527  	ws = memdb.NewWatchSet()
   528  	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "11")
   529  	if err != nil {
   530  		t.Fatalf("err: %v", err)
   531  	}
   532  
   533  	deploys = gatherDeploys(iter)
   534  	if len(deploys) != 2 {
   535  		t.Fatalf("err: %v", err)
   536  	}
   537  
   538  	iter, err = state.DeploymentsByIDPrefix(ws, deploy.Namespace, "1111")
   539  	if err != nil {
   540  		t.Fatalf("err: %v", err)
   541  	}
   542  
   543  	deploys = gatherDeploys(iter)
   544  	if len(deploys) != 1 {
   545  		t.Fatalf("err: %v", err)
   546  	}
   547  
   548  	if watchFired(ws) {
   549  		t.Fatalf("bad")
   550  	}
   551  }
   552  
   553  func TestStateStore_UpsertNode_Node(t *testing.T) {
   554  	require := require.New(t)
   555  	state := testStateStore(t)
   556  	node := mock.Node()
   557  
   558  	// Create a watchset so we can test that upsert fires the watch
   559  	ws := memdb.NewWatchSet()
   560  	_, err := state.NodeByID(ws, node.ID)
   561  	require.NoError(err)
   562  
   563  	require.NoError(state.UpsertNode(1000, node))
   564  	require.True(watchFired(ws))
   565  
   566  	ws = memdb.NewWatchSet()
   567  	out, err := state.NodeByID(ws, node.ID)
   568  	require.NoError(err)
   569  
   570  	out2, err := state.NodeBySecretID(ws, node.SecretID)
   571  	require.NoError(err)
   572  	require.EqualValues(node, out)
   573  	require.EqualValues(node, out2)
   574  	require.Len(out.Events, 1)
   575  	require.Equal(NodeRegisterEventRegistered, out.Events[0].Message)
   576  
   577  	index, err := state.Index("nodes")
   578  	require.NoError(err)
   579  	require.EqualValues(1000, index)
   580  	require.False(watchFired(ws))
   581  
   582  	// Transition the node to down and then up and ensure we get a re-register
   583  	// event
   584  	down := out.Copy()
   585  	down.Status = structs.NodeStatusDown
   586  	require.NoError(state.UpsertNode(1001, down))
   587  	require.NoError(state.UpsertNode(1002, out))
   588  
   589  	out, err = state.NodeByID(ws, node.ID)
   590  	require.NoError(err)
   591  	require.Len(out.Events, 2)
   592  	require.Equal(NodeRegisterEventReregistered, out.Events[1].Message)
   593  }
   594  
   595  func TestStateStore_DeleteNode_Node(t *testing.T) {
   596  	state := testStateStore(t)
   597  	node := mock.Node()
   598  
   599  	err := state.UpsertNode(1000, node)
   600  	if err != nil {
   601  		t.Fatalf("err: %v", err)
   602  	}
   603  
   604  	// Create a watchset so we can test that delete fires the watch
   605  	ws := memdb.NewWatchSet()
   606  	if _, err := state.NodeByID(ws, node.ID); err != nil {
   607  		t.Fatalf("bad: %v", err)
   608  	}
   609  
   610  	err = state.DeleteNode(1001, node.ID)
   611  	if err != nil {
   612  		t.Fatalf("err: %v", err)
   613  	}
   614  
   615  	if !watchFired(ws) {
   616  		t.Fatalf("bad")
   617  	}
   618  
   619  	ws = memdb.NewWatchSet()
   620  	out, err := state.NodeByID(ws, node.ID)
   621  	if err != nil {
   622  		t.Fatalf("err: %v", err)
   623  	}
   624  
   625  	if out != nil {
   626  		t.Fatalf("bad: %#v %#v", node, out)
   627  	}
   628  
   629  	index, err := state.Index("nodes")
   630  	if err != nil {
   631  		t.Fatalf("err: %v", err)
   632  	}
   633  	if index != 1001 {
   634  		t.Fatalf("bad: %d", index)
   635  	}
   636  
   637  	if watchFired(ws) {
   638  		t.Fatalf("bad")
   639  	}
   640  }
   641  
   642  func TestStateStore_UpdateNodeStatus_Node(t *testing.T) {
   643  	require := require.New(t)
   644  	state := testStateStore(t)
   645  	node := mock.Node()
   646  
   647  	require.NoError(state.UpsertNode(800, node))
   648  
   649  	// Create a watchset so we can test that update node status fires the watch
   650  	ws := memdb.NewWatchSet()
   651  	_, err := state.NodeByID(ws, node.ID)
   652  	require.NoError(err)
   653  
   654  	event := &structs.NodeEvent{
   655  		Message:   "Node ready foo",
   656  		Subsystem: structs.NodeEventSubsystemCluster,
   657  		Timestamp: time.Now(),
   658  	}
   659  
   660  	require.NoError(state.UpdateNodeStatus(801, node.ID, structs.NodeStatusReady, event))
   661  	require.True(watchFired(ws))
   662  
   663  	ws = memdb.NewWatchSet()
   664  	out, err := state.NodeByID(ws, node.ID)
   665  	require.NoError(err)
   666  	require.Equal(structs.NodeStatusReady, out.Status)
   667  	require.EqualValues(801, out.ModifyIndex)
   668  	require.Len(out.Events, 2)
   669  	require.Equal(event.Message, out.Events[1].Message)
   670  
   671  	index, err := state.Index("nodes")
   672  	require.NoError(err)
   673  	require.EqualValues(801, index)
   674  	require.False(watchFired(ws))
   675  }
   676  
   677  func TestStateStore_BatchUpdateNodeDrain(t *testing.T) {
   678  	require := require.New(t)
   679  	state := testStateStore(t)
   680  
   681  	n1, n2 := mock.Node(), mock.Node()
   682  	require.Nil(state.UpsertNode(1000, n1))
   683  	require.Nil(state.UpsertNode(1001, n2))
   684  
   685  	// Create a watchset so we can test that update node drain fires the watch
   686  	ws := memdb.NewWatchSet()
   687  	_, err := state.NodeByID(ws, n1.ID)
   688  	require.Nil(err)
   689  
   690  	expectedDrain := &structs.DrainStrategy{
   691  		DrainSpec: structs.DrainSpec{
   692  			Deadline: -1 * time.Second,
   693  		},
   694  	}
   695  
   696  	update := map[string]*structs.DrainUpdate{
   697  		n1.ID: {
   698  			DrainStrategy: expectedDrain,
   699  		},
   700  		n2.ID: {
   701  			DrainStrategy: expectedDrain,
   702  		},
   703  	}
   704  
   705  	event := &structs.NodeEvent{
   706  		Message:   "Drain strategy enabled",
   707  		Subsystem: structs.NodeEventSubsystemDrain,
   708  		Timestamp: time.Now(),
   709  	}
   710  	events := map[string]*structs.NodeEvent{
   711  		n1.ID: event,
   712  		n2.ID: event,
   713  	}
   714  
   715  	require.Nil(state.BatchUpdateNodeDrain(1002, update, events))
   716  	require.True(watchFired(ws))
   717  
   718  	ws = memdb.NewWatchSet()
   719  	for _, id := range []string{n1.ID, n2.ID} {
   720  		out, err := state.NodeByID(ws, id)
   721  		require.Nil(err)
   722  		require.True(out.Drain)
   723  		require.NotNil(out.DrainStrategy)
   724  		require.Equal(out.DrainStrategy, expectedDrain)
   725  		require.Len(out.Events, 2)
   726  		require.EqualValues(1002, out.ModifyIndex)
   727  	}
   728  
   729  	index, err := state.Index("nodes")
   730  	require.Nil(err)
   731  	require.EqualValues(1002, index)
   732  	require.False(watchFired(ws))
   733  }
   734  
   735  func TestStateStore_UpdateNodeDrain_Node(t *testing.T) {
   736  	require := require.New(t)
   737  	state := testStateStore(t)
   738  	node := mock.Node()
   739  
   740  	require.Nil(state.UpsertNode(1000, node))
   741  
   742  	// Create a watchset so we can test that update node drain fires the watch
   743  	ws := memdb.NewWatchSet()
   744  	_, err := state.NodeByID(ws, node.ID)
   745  	require.Nil(err)
   746  
   747  	expectedDrain := &structs.DrainStrategy{
   748  		DrainSpec: structs.DrainSpec{
   749  			Deadline: -1 * time.Second,
   750  		},
   751  	}
   752  
   753  	event := &structs.NodeEvent{
   754  		Message:   "Drain strategy enabled",
   755  		Subsystem: structs.NodeEventSubsystemDrain,
   756  		Timestamp: time.Now(),
   757  	}
   758  	require.Nil(state.UpdateNodeDrain(1001, node.ID, expectedDrain, false, event))
   759  	require.True(watchFired(ws))
   760  
   761  	ws = memdb.NewWatchSet()
   762  	out, err := state.NodeByID(ws, node.ID)
   763  	require.Nil(err)
   764  	require.True(out.Drain)
   765  	require.NotNil(out.DrainStrategy)
   766  	require.Equal(out.DrainStrategy, expectedDrain)
   767  	require.Len(out.Events, 2)
   768  	require.EqualValues(1001, out.ModifyIndex)
   769  
   770  	index, err := state.Index("nodes")
   771  	require.Nil(err)
   772  	require.EqualValues(1001, index)
   773  	require.False(watchFired(ws))
   774  }
   775  
   776  func TestStateStore_AddSingleNodeEvent(t *testing.T) {
   777  	require := require.New(t)
   778  	state := testStateStore(t)
   779  
   780  	node := mock.Node()
   781  
   782  	// We create a new node event every time we register a node
   783  	err := state.UpsertNode(1000, node)
   784  	require.Nil(err)
   785  
   786  	require.Equal(1, len(node.Events))
   787  	require.Equal(structs.NodeEventSubsystemCluster, node.Events[0].Subsystem)
   788  	require.Equal(NodeRegisterEventRegistered, node.Events[0].Message)
   789  
   790  	// Create a watchset so we can test that AddNodeEvent fires the watch
   791  	ws := memdb.NewWatchSet()
   792  	_, err = state.NodeByID(ws, node.ID)
   793  	require.Nil(err)
   794  
   795  	nodeEvent := &structs.NodeEvent{
   796  		Message:   "failed",
   797  		Subsystem: "Driver",
   798  		Timestamp: time.Now(),
   799  	}
   800  	nodeEvents := map[string][]*structs.NodeEvent{
   801  		node.ID: {nodeEvent},
   802  	}
   803  	err = state.UpsertNodeEvents(uint64(1001), nodeEvents)
   804  	require.Nil(err)
   805  
   806  	require.True(watchFired(ws))
   807  
   808  	ws = memdb.NewWatchSet()
   809  	out, err := state.NodeByID(ws, node.ID)
   810  	require.Nil(err)
   811  
   812  	require.Equal(2, len(out.Events))
   813  	require.Equal(nodeEvent, out.Events[1])
   814  }
   815  
   816  // To prevent stale node events from accumulating, we limit the number of
   817  // stored node events to 10.
   818  func TestStateStore_NodeEvents_RetentionWindow(t *testing.T) {
   819  	require := require.New(t)
   820  	state := testStateStore(t)
   821  
   822  	node := mock.Node()
   823  
   824  	err := state.UpsertNode(1000, node)
   825  	if err != nil {
   826  		t.Fatalf("err: %v", err)
   827  	}
   828  	require.Equal(1, len(node.Events))
   829  	require.Equal(structs.NodeEventSubsystemCluster, node.Events[0].Subsystem)
   830  	require.Equal(NodeRegisterEventRegistered, node.Events[0].Message)
   831  
   832  	var out *structs.Node
   833  	for i := 1; i <= 20; i++ {
   834  		ws := memdb.NewWatchSet()
   835  		out, err = state.NodeByID(ws, node.ID)
   836  		require.Nil(err)
   837  
   838  		nodeEvent := &structs.NodeEvent{
   839  			Message:   fmt.Sprintf("%dith failed", i),
   840  			Subsystem: "Driver",
   841  			Timestamp: time.Now(),
   842  		}
   843  
   844  		nodeEvents := map[string][]*structs.NodeEvent{
   845  			out.ID: {nodeEvent},
   846  		}
   847  		err := state.UpsertNodeEvents(uint64(i), nodeEvents)
   848  		require.Nil(err)
   849  
   850  		require.True(watchFired(ws))
   851  		ws = memdb.NewWatchSet()
   852  		out, err = state.NodeByID(ws, node.ID)
   853  		require.Nil(err)
   854  	}
   855  
   856  	ws := memdb.NewWatchSet()
   857  	out, err = state.NodeByID(ws, node.ID)
   858  	require.Nil(err)
   859  
   860  	require.Equal(10, len(out.Events))
   861  	require.Equal(uint64(11), out.Events[0].CreateIndex)
   862  	require.Equal(uint64(20), out.Events[len(out.Events)-1].CreateIndex)
   863  }
   864  
   865  func TestStateStore_UpdateNodeDrain_ResetEligiblity(t *testing.T) {
   866  	require := require.New(t)
   867  	state := testStateStore(t)
   868  	node := mock.Node()
   869  	require.Nil(state.UpsertNode(1000, node))
   870  
   871  	// Create a watchset so we can test that update node drain fires the watch
   872  	ws := memdb.NewWatchSet()
   873  	_, err := state.NodeByID(ws, node.ID)
   874  	require.Nil(err)
   875  
   876  	drain := &structs.DrainStrategy{
   877  		DrainSpec: structs.DrainSpec{
   878  			Deadline: -1 * time.Second,
   879  		},
   880  	}
   881  
   882  	event1 := &structs.NodeEvent{
   883  		Message:   "Drain strategy enabled",
   884  		Subsystem: structs.NodeEventSubsystemDrain,
   885  		Timestamp: time.Now(),
   886  	}
   887  	require.Nil(state.UpdateNodeDrain(1001, node.ID, drain, false, event1))
   888  	require.True(watchFired(ws))
   889  
   890  	// Remove the drain
   891  	event2 := &structs.NodeEvent{
   892  		Message:   "Drain strategy disabled",
   893  		Subsystem: structs.NodeEventSubsystemDrain,
   894  		Timestamp: time.Now(),
   895  	}
   896  	require.Nil(state.UpdateNodeDrain(1002, node.ID, nil, true, event2))
   897  
   898  	ws = memdb.NewWatchSet()
   899  	out, err := state.NodeByID(ws, node.ID)
   900  	require.Nil(err)
   901  	require.False(out.Drain)
   902  	require.Nil(out.DrainStrategy)
   903  	require.Equal(out.SchedulingEligibility, structs.NodeSchedulingEligible)
   904  	require.Len(out.Events, 3)
   905  	require.EqualValues(1002, out.ModifyIndex)
   906  
   907  	index, err := state.Index("nodes")
   908  	require.Nil(err)
   909  	require.EqualValues(1002, index)
   910  	require.False(watchFired(ws))
   911  }
   912  
   913  func TestStateStore_UpdateNodeEligibility(t *testing.T) {
   914  	require := require.New(t)
   915  	state := testStateStore(t)
   916  	node := mock.Node()
   917  
   918  	err := state.UpsertNode(1000, node)
   919  	if err != nil {
   920  		t.Fatalf("err: %v", err)
   921  	}
   922  
   923  	expectedEligibility := structs.NodeSchedulingIneligible
   924  
   925  	// Create a watchset so we can test that update node drain fires the watch
   926  	ws := memdb.NewWatchSet()
   927  	if _, err := state.NodeByID(ws, node.ID); err != nil {
   928  		t.Fatalf("bad: %v", err)
   929  	}
   930  
   931  	event := &structs.NodeEvent{
   932  		Message:   "Node marked as ineligible",
   933  		Subsystem: structs.NodeEventSubsystemCluster,
   934  		Timestamp: time.Now(),
   935  	}
   936  	require.Nil(state.UpdateNodeEligibility(1001, node.ID, expectedEligibility, event))
   937  	require.True(watchFired(ws))
   938  
   939  	ws = memdb.NewWatchSet()
   940  	out, err := state.NodeByID(ws, node.ID)
   941  	require.Nil(err)
   942  	require.Equal(out.SchedulingEligibility, expectedEligibility)
   943  	require.Len(out.Events, 2)
   944  	require.Equal(out.Events[1], event)
   945  	require.EqualValues(1001, out.ModifyIndex)
   946  
   947  	index, err := state.Index("nodes")
   948  	require.Nil(err)
   949  	require.EqualValues(1001, index)
   950  	require.False(watchFired(ws))
   951  
   952  	// Set a drain strategy
   953  	expectedDrain := &structs.DrainStrategy{
   954  		DrainSpec: structs.DrainSpec{
   955  			Deadline: -1 * time.Second,
   956  		},
   957  	}
   958  	require.Nil(state.UpdateNodeDrain(1002, node.ID, expectedDrain, false, nil))
   959  
   960  	// Try to set the node to eligible
   961  	err = state.UpdateNodeEligibility(1003, node.ID, structs.NodeSchedulingEligible, nil)
   962  	require.NotNil(err)
   963  	require.Contains(err.Error(), "while it is draining")
   964  }
   965  
   966  func TestStateStore_Nodes(t *testing.T) {
   967  	state := testStateStore(t)
   968  	var nodes []*structs.Node
   969  
   970  	for i := 0; i < 10; i++ {
   971  		node := mock.Node()
   972  		nodes = append(nodes, node)
   973  
   974  		err := state.UpsertNode(1000+uint64(i), node)
   975  		if err != nil {
   976  			t.Fatalf("err: %v", err)
   977  		}
   978  	}
   979  
   980  	// Create a watchset so we can test that getters don't cause it to fire
   981  	ws := memdb.NewWatchSet()
   982  	iter, err := state.Nodes(ws)
   983  	if err != nil {
   984  		t.Fatalf("bad: %v", err)
   985  	}
   986  
   987  	var out []*structs.Node
   988  	for {
   989  		raw := iter.Next()
   990  		if raw == nil {
   991  			break
   992  		}
   993  		out = append(out, raw.(*structs.Node))
   994  	}
   995  
   996  	sort.Sort(NodeIDSort(nodes))
   997  	sort.Sort(NodeIDSort(out))
   998  
   999  	if !reflect.DeepEqual(nodes, out) {
  1000  		t.Fatalf("bad: %#v %#v", nodes, out)
  1001  	}
  1002  
  1003  	if watchFired(ws) {
  1004  		t.Fatalf("bad")
  1005  	}
  1006  }
  1007  
  1008  func TestStateStore_NodesByIDPrefix(t *testing.T) {
  1009  	state := testStateStore(t)
  1010  	node := mock.Node()
  1011  
  1012  	node.ID = "11111111-662e-d0ab-d1c9-3e434af7bdb4"
  1013  	err := state.UpsertNode(1000, node)
  1014  	if err != nil {
  1015  		t.Fatalf("err: %v", err)
  1016  	}
  1017  
  1018  	// Create a watchset so we can test that getters don't cause it to fire
  1019  	ws := memdb.NewWatchSet()
  1020  	iter, err := state.NodesByIDPrefix(ws, node.ID)
  1021  	if err != nil {
  1022  		t.Fatalf("err: %v", err)
  1023  	}
  1024  
  1025  	gatherNodes := func(iter memdb.ResultIterator) []*structs.Node {
  1026  		var nodes []*structs.Node
  1027  		for {
  1028  			raw := iter.Next()
  1029  			if raw == nil {
  1030  				break
  1031  			}
  1032  			node := raw.(*structs.Node)
  1033  			nodes = append(nodes, node)
  1034  		}
  1035  		return nodes
  1036  	}
  1037  
  1038  	nodes := gatherNodes(iter)
  1039  	if len(nodes) != 1 {
  1040  		t.Fatalf("err: %v", err)
  1041  	}
  1042  
  1043  	if watchFired(ws) {
  1044  		t.Fatalf("bad")
  1045  	}
  1046  
  1047  	iter, err = state.NodesByIDPrefix(ws, "11")
  1048  	if err != nil {
  1049  		t.Fatalf("err: %v", err)
  1050  	}
  1051  
  1052  	nodes = gatherNodes(iter)
  1053  	if len(nodes) != 1 {
  1054  		t.Fatalf("err: %v", err)
  1055  	}
  1056  
  1057  	node = mock.Node()
  1058  	node.ID = "11222222-662e-d0ab-d1c9-3e434af7bdb4"
  1059  	err = state.UpsertNode(1001, node)
  1060  	if err != nil {
  1061  		t.Fatalf("err: %v", err)
  1062  	}
  1063  
  1064  	if !watchFired(ws) {
  1065  		t.Fatalf("bad")
  1066  	}
  1067  
  1068  	ws = memdb.NewWatchSet()
  1069  	iter, err = state.NodesByIDPrefix(ws, "11")
  1070  	if err != nil {
  1071  		t.Fatalf("err: %v", err)
  1072  	}
  1073  
  1074  	nodes = gatherNodes(iter)
  1075  	if len(nodes) != 2 {
  1076  		t.Fatalf("err: %v", err)
  1077  	}
  1078  
  1079  	iter, err = state.NodesByIDPrefix(ws, "1111")
  1080  	if err != nil {
  1081  		t.Fatalf("err: %v", err)
  1082  	}
  1083  
  1084  	nodes = gatherNodes(iter)
  1085  	if len(nodes) != 1 {
  1086  		t.Fatalf("err: %v", err)
  1087  	}
  1088  
  1089  	if watchFired(ws) {
  1090  		t.Fatalf("bad")
  1091  	}
  1092  }
  1093  
  1094  func TestStateStore_RestoreNode(t *testing.T) {
  1095  	state := testStateStore(t)
  1096  	node := mock.Node()
  1097  
  1098  	restore, err := state.Restore()
  1099  	if err != nil {
  1100  		t.Fatalf("err: %v", err)
  1101  	}
  1102  
  1103  	err = restore.NodeRestore(node)
  1104  	if err != nil {
  1105  		t.Fatalf("err: %v", err)
  1106  	}
  1107  	restore.Commit()
  1108  
  1109  	ws := memdb.NewWatchSet()
  1110  	out, err := state.NodeByID(ws, node.ID)
  1111  	if err != nil {
  1112  		t.Fatalf("err: %v", err)
  1113  	}
  1114  
  1115  	if !reflect.DeepEqual(out, node) {
  1116  		t.Fatalf("Bad: %#v %#v", out, node)
  1117  	}
  1118  }
  1119  
  1120  func TestStateStore_UpsertJob_Job(t *testing.T) {
  1121  	state := testStateStore(t)
  1122  	job := mock.Job()
  1123  
  1124  	// Create a watchset so we can test that upsert fires the watch
  1125  	ws := memdb.NewWatchSet()
  1126  	_, err := state.JobByID(ws, job.Namespace, job.ID)
  1127  	if err != nil {
  1128  		t.Fatalf("bad: %v", err)
  1129  	}
  1130  
  1131  	if err := state.UpsertJob(1000, job); err != nil {
  1132  		t.Fatalf("err: %v", err)
  1133  	}
  1134  	if !watchFired(ws) {
  1135  		t.Fatalf("bad")
  1136  	}
  1137  
  1138  	ws = memdb.NewWatchSet()
  1139  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1140  	if err != nil {
  1141  		t.Fatalf("err: %v", err)
  1142  	}
  1143  
  1144  	if !reflect.DeepEqual(job, out) {
  1145  		t.Fatalf("bad: %#v %#v", job, out)
  1146  	}
  1147  
  1148  	index, err := state.Index("jobs")
  1149  	if err != nil {
  1150  		t.Fatalf("err: %v", err)
  1151  	}
  1152  	if index != 1000 {
  1153  		t.Fatalf("bad: %d", index)
  1154  	}
  1155  
  1156  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1157  	if err != nil {
  1158  		t.Fatalf("err: %v", err)
  1159  	}
  1160  	if summary == nil {
  1161  		t.Fatalf("nil summary")
  1162  	}
  1163  	if summary.JobID != job.ID {
  1164  		t.Fatalf("bad summary id: %v", summary.JobID)
  1165  	}
  1166  	_, ok := summary.Summary["web"]
  1167  	if !ok {
  1168  		t.Fatalf("nil summary for task group")
  1169  	}
  1170  	if watchFired(ws) {
  1171  		t.Fatalf("bad")
  1172  	}
  1173  
  1174  	// Check the job versions
  1175  	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1176  	if err != nil {
  1177  		t.Fatalf("err: %v", err)
  1178  	}
  1179  	if len(allVersions) != 1 {
  1180  		t.Fatalf("got %d; want 1", len(allVersions))
  1181  	}
  1182  
  1183  	if a := allVersions[0]; a.ID != job.ID || a.Version != 0 {
  1184  		t.Fatalf("bad: %v", a)
  1185  	}
  1186  
  1187  	// Test the looking up the job by version returns the same results
  1188  	vout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
  1189  	if err != nil {
  1190  		t.Fatalf("err: %v", err)
  1191  	}
  1192  
  1193  	if !reflect.DeepEqual(out, vout) {
  1194  		t.Fatalf("bad: %#v %#v", out, vout)
  1195  	}
  1196  }
  1197  
  1198  func TestStateStore_UpdateUpsertJob_Job(t *testing.T) {
  1199  	state := testStateStore(t)
  1200  	job := mock.Job()
  1201  
  1202  	// Create a watchset so we can test that upsert fires the watch
  1203  	ws := memdb.NewWatchSet()
  1204  	_, err := state.JobByID(ws, job.Namespace, job.ID)
  1205  	if err != nil {
  1206  		t.Fatalf("bad: %v", err)
  1207  	}
  1208  
  1209  	if err := state.UpsertJob(1000, job); err != nil {
  1210  		t.Fatalf("err: %v", err)
  1211  	}
  1212  
  1213  	job2 := mock.Job()
  1214  	job2.ID = job.ID
  1215  	job2.AllAtOnce = true
  1216  	err = state.UpsertJob(1001, job2)
  1217  	if err != nil {
  1218  		t.Fatalf("err: %v", err)
  1219  	}
  1220  
  1221  	if !watchFired(ws) {
  1222  		t.Fatalf("bad")
  1223  	}
  1224  
  1225  	ws = memdb.NewWatchSet()
  1226  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1227  	if err != nil {
  1228  		t.Fatalf("err: %v", err)
  1229  	}
  1230  
  1231  	if !reflect.DeepEqual(job2, out) {
  1232  		t.Fatalf("bad: %#v %#v", job2, out)
  1233  	}
  1234  
  1235  	if out.CreateIndex != 1000 {
  1236  		t.Fatalf("bad: %#v", out)
  1237  	}
  1238  	if out.ModifyIndex != 1001 {
  1239  		t.Fatalf("bad: %#v", out)
  1240  	}
  1241  	if out.Version != 1 {
  1242  		t.Fatalf("bad: %#v", out)
  1243  	}
  1244  
  1245  	index, err := state.Index("jobs")
  1246  	if err != nil {
  1247  		t.Fatalf("err: %v", err)
  1248  	}
  1249  	if index != 1001 {
  1250  		t.Fatalf("bad: %d", index)
  1251  	}
  1252  
  1253  	// Test the looking up the job by version returns the same results
  1254  	vout, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 1)
  1255  	if err != nil {
  1256  		t.Fatalf("err: %v", err)
  1257  	}
  1258  
  1259  	if !reflect.DeepEqual(out, vout) {
  1260  		t.Fatalf("bad: %#v %#v", out, vout)
  1261  	}
  1262  
  1263  	// Test that the job summary remains the same if the job is updated but
  1264  	// count remains same
  1265  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1266  	if err != nil {
  1267  		t.Fatalf("err: %v", err)
  1268  	}
  1269  	if summary == nil {
  1270  		t.Fatalf("nil summary")
  1271  	}
  1272  	if summary.JobID != job.ID {
  1273  		t.Fatalf("bad summary id: %v", summary.JobID)
  1274  	}
  1275  	_, ok := summary.Summary["web"]
  1276  	if !ok {
  1277  		t.Fatalf("nil summary for task group")
  1278  	}
  1279  
  1280  	// Check the job versions
  1281  	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1282  	if err != nil {
  1283  		t.Fatalf("err: %v", err)
  1284  	}
  1285  	if len(allVersions) != 2 {
  1286  		t.Fatalf("got %d; want 1", len(allVersions))
  1287  	}
  1288  
  1289  	if a := allVersions[0]; a.ID != job.ID || a.Version != 1 || !a.AllAtOnce {
  1290  		t.Fatalf("bad: %+v", a)
  1291  	}
  1292  	if a := allVersions[1]; a.ID != job.ID || a.Version != 0 || a.AllAtOnce {
  1293  		t.Fatalf("bad: %+v", a)
  1294  	}
  1295  
  1296  	if watchFired(ws) {
  1297  		t.Fatalf("bad")
  1298  	}
  1299  }
  1300  
  1301  func TestStateStore_UpdateUpsertJob_PeriodicJob(t *testing.T) {
  1302  	state := testStateStore(t)
  1303  	job := mock.PeriodicJob()
  1304  
  1305  	// Create a watchset so we can test that upsert fires the watch
  1306  	ws := memdb.NewWatchSet()
  1307  	_, err := state.JobByID(ws, job.Namespace, job.ID)
  1308  	if err != nil {
  1309  		t.Fatalf("bad: %v", err)
  1310  	}
  1311  
  1312  	if err := state.UpsertJob(1000, job); err != nil {
  1313  		t.Fatalf("err: %v", err)
  1314  	}
  1315  
  1316  	// Create a child and an evaluation
  1317  	job2 := job.Copy()
  1318  	job2.Periodic = nil
  1319  	job2.ID = fmt.Sprintf("%v/%s-1490635020", job.ID, structs.PeriodicLaunchSuffix)
  1320  	err = state.UpsertJob(1001, job2)
  1321  	if err != nil {
  1322  		t.Fatalf("err: %v", err)
  1323  	}
  1324  
  1325  	eval := mock.Eval()
  1326  	eval.JobID = job2.ID
  1327  	err = state.UpsertEvals(1002, []*structs.Evaluation{eval})
  1328  	if err != nil {
  1329  		t.Fatalf("err: %v", err)
  1330  	}
  1331  
  1332  	job3 := job.Copy()
  1333  	job3.TaskGroups[0].Tasks[0].Name = "new name"
  1334  	err = state.UpsertJob(1003, job3)
  1335  	if err != nil {
  1336  		t.Fatalf("err: %v", err)
  1337  	}
  1338  
  1339  	if !watchFired(ws) {
  1340  		t.Fatalf("bad")
  1341  	}
  1342  
  1343  	ws = memdb.NewWatchSet()
  1344  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1345  	if err != nil {
  1346  		t.Fatalf("err: %v", err)
  1347  	}
  1348  
  1349  	if s, e := out.Status, structs.JobStatusRunning; s != e {
  1350  		t.Fatalf("got status %v; want %v", s, e)
  1351  	}
  1352  
  1353  }
  1354  
  1355  // This test ensures that UpsertJob creates the EphemeralDisk is a job doesn't have
  1356  // one and clear out the task's disk resource asks
  1357  // COMPAT 0.4.1 -> 0.5
  1358  func TestStateStore_UpsertJob_NoEphemeralDisk(t *testing.T) {
  1359  	state := testStateStore(t)
  1360  	job := mock.Job()
  1361  
  1362  	// Set the EphemeralDisk to nil and set the tasks's DiskMB to 150
  1363  	job.TaskGroups[0].EphemeralDisk = nil
  1364  	job.TaskGroups[0].Tasks[0].Resources.DiskMB = 150
  1365  
  1366  	err := state.UpsertJob(1000, job)
  1367  	if err != nil {
  1368  		t.Fatalf("err: %v", err)
  1369  	}
  1370  
  1371  	ws := memdb.NewWatchSet()
  1372  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1373  	if err != nil {
  1374  		t.Fatalf("err: %v", err)
  1375  	}
  1376  
  1377  	// Expect the state store to create the EphemeralDisk and clear out Tasks's
  1378  	// DiskMB
  1379  	expected := job.Copy()
  1380  	expected.TaskGroups[0].EphemeralDisk = &structs.EphemeralDisk{
  1381  		SizeMB: 150,
  1382  	}
  1383  	expected.TaskGroups[0].Tasks[0].Resources.DiskMB = 0
  1384  
  1385  	if !reflect.DeepEqual(expected, out) {
  1386  		t.Fatalf("bad: %#v %#v", expected, out)
  1387  	}
  1388  }
  1389  
  1390  func TestStateStore_UpsertJob_BadNamespace(t *testing.T) {
  1391  	assert := assert.New(t)
  1392  	state := testStateStore(t)
  1393  	job := mock.Job()
  1394  	job.Namespace = "foo"
  1395  
  1396  	err := state.UpsertJob(1000, job)
  1397  	assert.Contains(err.Error(), "nonexistent namespace")
  1398  
  1399  	ws := memdb.NewWatchSet()
  1400  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1401  	assert.Nil(err)
  1402  	assert.Nil(out)
  1403  }
  1404  
  1405  // Upsert a job that is the child of a parent job and ensures its summary gets
  1406  // updated.
  1407  func TestStateStore_UpsertJob_ChildJob(t *testing.T) {
  1408  	state := testStateStore(t)
  1409  
  1410  	// Create a watchset so we can test that upsert fires the watch
  1411  	parent := mock.Job()
  1412  	ws := memdb.NewWatchSet()
  1413  	_, err := state.JobByID(ws, parent.Namespace, parent.ID)
  1414  	if err != nil {
  1415  		t.Fatalf("bad: %v", err)
  1416  	}
  1417  
  1418  	if err := state.UpsertJob(1000, parent); err != nil {
  1419  		t.Fatalf("err: %v", err)
  1420  	}
  1421  
  1422  	child := mock.Job()
  1423  	child.ParentID = parent.ID
  1424  	if err := state.UpsertJob(1001, child); err != nil {
  1425  		t.Fatalf("err: %v", err)
  1426  	}
  1427  
  1428  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  1429  	if err != nil {
  1430  		t.Fatalf("err: %v", err)
  1431  	}
  1432  	if summary == nil {
  1433  		t.Fatalf("nil summary")
  1434  	}
  1435  	if summary.JobID != parent.ID {
  1436  		t.Fatalf("bad summary id: %v", parent.ID)
  1437  	}
  1438  	if summary.Children == nil {
  1439  		t.Fatalf("nil children summary")
  1440  	}
  1441  	if summary.Children.Pending != 1 || summary.Children.Running != 0 || summary.Children.Dead != 0 {
  1442  		t.Fatalf("bad children summary: %v", summary.Children)
  1443  	}
  1444  	if !watchFired(ws) {
  1445  		t.Fatalf("bad")
  1446  	}
  1447  }
  1448  
  1449  func TestStateStore_UpdateUpsertJob_JobVersion(t *testing.T) {
  1450  	state := testStateStore(t)
  1451  
  1452  	// Create a job and mark it as stable
  1453  	job := mock.Job()
  1454  	job.Stable = true
  1455  	job.Name = "0"
  1456  
  1457  	// Create a watchset so we can test that upsert fires the watch
  1458  	ws := memdb.NewWatchSet()
  1459  	_, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1460  	if err != nil {
  1461  		t.Fatalf("bad: %v", err)
  1462  	}
  1463  
  1464  	if err := state.UpsertJob(1000, job); err != nil {
  1465  		t.Fatalf("err: %v", err)
  1466  	}
  1467  
  1468  	if !watchFired(ws) {
  1469  		t.Fatalf("bad")
  1470  	}
  1471  
  1472  	var finalJob *structs.Job
  1473  	for i := 1; i < 300; i++ {
  1474  		finalJob = mock.Job()
  1475  		finalJob.ID = job.ID
  1476  		finalJob.Name = fmt.Sprintf("%d", i)
  1477  		err = state.UpsertJob(uint64(1000+i), finalJob)
  1478  		if err != nil {
  1479  			t.Fatalf("err: %v", err)
  1480  		}
  1481  	}
  1482  
  1483  	ws = memdb.NewWatchSet()
  1484  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1485  	if err != nil {
  1486  		t.Fatalf("err: %v", err)
  1487  	}
  1488  
  1489  	if !reflect.DeepEqual(finalJob, out) {
  1490  		t.Fatalf("bad: %#v %#v", finalJob, out)
  1491  	}
  1492  
  1493  	if out.CreateIndex != 1000 {
  1494  		t.Fatalf("bad: %#v", out)
  1495  	}
  1496  	if out.ModifyIndex != 1299 {
  1497  		t.Fatalf("bad: %#v", out)
  1498  	}
  1499  	if out.Version != 299 {
  1500  		t.Fatalf("bad: %#v", out)
  1501  	}
  1502  
  1503  	index, err := state.Index("job_version")
  1504  	if err != nil {
  1505  		t.Fatalf("err: %v", err)
  1506  	}
  1507  	if index != 1299 {
  1508  		t.Fatalf("bad: %d", index)
  1509  	}
  1510  
  1511  	// Check the job versions
  1512  	allVersions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1513  	if err != nil {
  1514  		t.Fatalf("err: %v", err)
  1515  	}
  1516  	if len(allVersions) != structs.JobTrackedVersions {
  1517  		t.Fatalf("got %d; want %d", len(allVersions), structs.JobTrackedVersions)
  1518  	}
  1519  
  1520  	if a := allVersions[0]; a.ID != job.ID || a.Version != 299 || a.Name != "299" {
  1521  		t.Fatalf("bad: %+v", a)
  1522  	}
  1523  	if a := allVersions[1]; a.ID != job.ID || a.Version != 298 || a.Name != "298" {
  1524  		t.Fatalf("bad: %+v", a)
  1525  	}
  1526  
  1527  	// Ensure we didn't delete the stable job
  1528  	if a := allVersions[structs.JobTrackedVersions-1]; a.ID != job.ID ||
  1529  		a.Version != 0 || a.Name != "0" || !a.Stable {
  1530  		t.Fatalf("bad: %+v", a)
  1531  	}
  1532  
  1533  	if watchFired(ws) {
  1534  		t.Fatalf("bad")
  1535  	}
  1536  }
  1537  
  1538  func TestStateStore_DeleteJob_Job(t *testing.T) {
  1539  	state := testStateStore(t)
  1540  	job := mock.Job()
  1541  
  1542  	err := state.UpsertJob(1000, job)
  1543  	if err != nil {
  1544  		t.Fatalf("err: %v", err)
  1545  	}
  1546  
  1547  	// Create a watchset so we can test that delete fires the watch
  1548  	ws := memdb.NewWatchSet()
  1549  	if _, err := state.JobByID(ws, job.Namespace, job.ID); err != nil {
  1550  		t.Fatalf("bad: %v", err)
  1551  	}
  1552  
  1553  	err = state.DeleteJob(1001, job.Namespace, job.ID)
  1554  	if err != nil {
  1555  		t.Fatalf("err: %v", err)
  1556  	}
  1557  
  1558  	if !watchFired(ws) {
  1559  		t.Fatalf("bad")
  1560  	}
  1561  
  1562  	ws = memdb.NewWatchSet()
  1563  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1564  	if err != nil {
  1565  		t.Fatalf("err: %v", err)
  1566  	}
  1567  
  1568  	if out != nil {
  1569  		t.Fatalf("bad: %#v %#v", job, out)
  1570  	}
  1571  
  1572  	index, err := state.Index("jobs")
  1573  	if err != nil {
  1574  		t.Fatalf("err: %v", err)
  1575  	}
  1576  	if index != 1001 {
  1577  		t.Fatalf("bad: %d", index)
  1578  	}
  1579  
  1580  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1581  	if err != nil {
  1582  		t.Fatalf("err: %v", err)
  1583  	}
  1584  	if summary != nil {
  1585  		t.Fatalf("expected summary to be nil, but got: %v", summary)
  1586  	}
  1587  
  1588  	index, err = state.Index("job_summary")
  1589  	if err != nil {
  1590  		t.Fatalf("err: %v", err)
  1591  	}
  1592  	if index != 1001 {
  1593  		t.Fatalf("bad: %d", index)
  1594  	}
  1595  
  1596  	versions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1597  	if err != nil {
  1598  		t.Fatalf("err: %v", err)
  1599  	}
  1600  	if len(versions) != 0 {
  1601  		t.Fatalf("expected no job versions")
  1602  	}
  1603  
  1604  	index, err = state.Index("job_summary")
  1605  	if err != nil {
  1606  		t.Fatalf("err: %v", err)
  1607  	}
  1608  	if index != 1001 {
  1609  		t.Fatalf("bad: %d", index)
  1610  	}
  1611  
  1612  	if watchFired(ws) {
  1613  		t.Fatalf("bad")
  1614  	}
  1615  }
  1616  
  1617  func TestStateStore_DeleteJob_MultipleVersions(t *testing.T) {
  1618  	state := testStateStore(t)
  1619  	assert := assert.New(t)
  1620  
  1621  	// Create a job and mark it as stable
  1622  	job := mock.Job()
  1623  	job.Stable = true
  1624  	job.Priority = 0
  1625  
  1626  	// Create a watchset so we can test that upsert fires the watch
  1627  	ws := memdb.NewWatchSet()
  1628  	_, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1629  	assert.Nil(err)
  1630  	assert.Nil(state.UpsertJob(1000, job))
  1631  	assert.True(watchFired(ws))
  1632  
  1633  	var finalJob *structs.Job
  1634  	for i := 1; i < 20; i++ {
  1635  		finalJob = mock.Job()
  1636  		finalJob.ID = job.ID
  1637  		finalJob.Priority = i
  1638  		assert.Nil(state.UpsertJob(uint64(1000+i), finalJob))
  1639  	}
  1640  
  1641  	assert.Nil(state.DeleteJob(1020, job.Namespace, job.ID))
  1642  	assert.True(watchFired(ws))
  1643  
  1644  	ws = memdb.NewWatchSet()
  1645  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  1646  	assert.Nil(err)
  1647  	assert.Nil(out)
  1648  
  1649  	index, err := state.Index("jobs")
  1650  	assert.Nil(err)
  1651  	assert.EqualValues(1020, index)
  1652  
  1653  	summary, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  1654  	assert.Nil(err)
  1655  	assert.Nil(summary)
  1656  
  1657  	index, err = state.Index("job_version")
  1658  	assert.Nil(err)
  1659  	assert.EqualValues(1020, index)
  1660  
  1661  	versions, err := state.JobVersionsByID(ws, job.Namespace, job.ID)
  1662  	assert.Nil(err)
  1663  	assert.Len(versions, 0)
  1664  
  1665  	index, err = state.Index("job_summary")
  1666  	assert.Nil(err)
  1667  	assert.EqualValues(1020, index)
  1668  
  1669  	assert.False(watchFired(ws))
  1670  }
  1671  
  1672  func TestStateStore_DeleteJob_ChildJob(t *testing.T) {
  1673  	state := testStateStore(t)
  1674  
  1675  	parent := mock.Job()
  1676  	if err := state.UpsertJob(998, parent); err != nil {
  1677  		t.Fatalf("err: %v", err)
  1678  	}
  1679  
  1680  	child := mock.Job()
  1681  	child.ParentID = parent.ID
  1682  
  1683  	if err := state.UpsertJob(999, child); err != nil {
  1684  		t.Fatalf("err: %v", err)
  1685  	}
  1686  
  1687  	// Create a watchset so we can test that delete fires the watch
  1688  	ws := memdb.NewWatchSet()
  1689  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  1690  		t.Fatalf("bad: %v", err)
  1691  	}
  1692  
  1693  	err := state.DeleteJob(1001, child.Namespace, child.ID)
  1694  	if err != nil {
  1695  		t.Fatalf("err: %v", err)
  1696  	}
  1697  	if !watchFired(ws) {
  1698  		t.Fatalf("bad")
  1699  	}
  1700  
  1701  	ws = memdb.NewWatchSet()
  1702  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  1703  	if err != nil {
  1704  		t.Fatalf("err: %v", err)
  1705  	}
  1706  	if summary == nil {
  1707  		t.Fatalf("nil summary")
  1708  	}
  1709  	if summary.JobID != parent.ID {
  1710  		t.Fatalf("bad summary id: %v", parent.ID)
  1711  	}
  1712  	if summary.Children == nil {
  1713  		t.Fatalf("nil children summary")
  1714  	}
  1715  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  1716  		t.Fatalf("bad children summary: %v", summary.Children)
  1717  	}
  1718  	if watchFired(ws) {
  1719  		t.Fatalf("bad")
  1720  	}
  1721  }
  1722  
  1723  func TestStateStore_Jobs(t *testing.T) {
  1724  	state := testStateStore(t)
  1725  	var jobs []*structs.Job
  1726  
  1727  	for i := 0; i < 10; i++ {
  1728  		job := mock.Job()
  1729  		jobs = append(jobs, job)
  1730  
  1731  		err := state.UpsertJob(1000+uint64(i), job)
  1732  		if err != nil {
  1733  			t.Fatalf("err: %v", err)
  1734  		}
  1735  	}
  1736  
  1737  	ws := memdb.NewWatchSet()
  1738  	iter, err := state.Jobs(ws)
  1739  	if err != nil {
  1740  		t.Fatalf("err: %v", err)
  1741  	}
  1742  
  1743  	var out []*structs.Job
  1744  	for {
  1745  		raw := iter.Next()
  1746  		if raw == nil {
  1747  			break
  1748  		}
  1749  		out = append(out, raw.(*structs.Job))
  1750  	}
  1751  
  1752  	sort.Sort(JobIDSort(jobs))
  1753  	sort.Sort(JobIDSort(out))
  1754  
  1755  	if !reflect.DeepEqual(jobs, out) {
  1756  		t.Fatalf("bad: %#v %#v", jobs, out)
  1757  	}
  1758  	if watchFired(ws) {
  1759  		t.Fatalf("bad")
  1760  	}
  1761  }
  1762  
  1763  func TestStateStore_JobVersions(t *testing.T) {
  1764  	state := testStateStore(t)
  1765  	var jobs []*structs.Job
  1766  
  1767  	for i := 0; i < 10; i++ {
  1768  		job := mock.Job()
  1769  		jobs = append(jobs, job)
  1770  
  1771  		err := state.UpsertJob(1000+uint64(i), job)
  1772  		if err != nil {
  1773  			t.Fatalf("err: %v", err)
  1774  		}
  1775  	}
  1776  
  1777  	ws := memdb.NewWatchSet()
  1778  	iter, err := state.JobVersions(ws)
  1779  	if err != nil {
  1780  		t.Fatalf("err: %v", err)
  1781  	}
  1782  
  1783  	var out []*structs.Job
  1784  	for {
  1785  		raw := iter.Next()
  1786  		if raw == nil {
  1787  			break
  1788  		}
  1789  		out = append(out, raw.(*structs.Job))
  1790  	}
  1791  
  1792  	sort.Sort(JobIDSort(jobs))
  1793  	sort.Sort(JobIDSort(out))
  1794  
  1795  	if !reflect.DeepEqual(jobs, out) {
  1796  		t.Fatalf("bad: %#v %#v", jobs, out)
  1797  	}
  1798  	if watchFired(ws) {
  1799  		t.Fatalf("bad")
  1800  	}
  1801  }
  1802  
  1803  func TestStateStore_JobsByIDPrefix(t *testing.T) {
  1804  	state := testStateStore(t)
  1805  	job := mock.Job()
  1806  
  1807  	job.ID = "redis"
  1808  	err := state.UpsertJob(1000, job)
  1809  	if err != nil {
  1810  		t.Fatalf("err: %v", err)
  1811  	}
  1812  
  1813  	ws := memdb.NewWatchSet()
  1814  	iter, err := state.JobsByIDPrefix(ws, job.Namespace, job.ID)
  1815  	if err != nil {
  1816  		t.Fatalf("err: %v", err)
  1817  	}
  1818  
  1819  	gatherJobs := func(iter memdb.ResultIterator) []*structs.Job {
  1820  		var jobs []*structs.Job
  1821  		for {
  1822  			raw := iter.Next()
  1823  			if raw == nil {
  1824  				break
  1825  			}
  1826  			jobs = append(jobs, raw.(*structs.Job))
  1827  		}
  1828  		return jobs
  1829  	}
  1830  
  1831  	jobs := gatherJobs(iter)
  1832  	if len(jobs) != 1 {
  1833  		t.Fatalf("err: %v", err)
  1834  	}
  1835  
  1836  	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "re")
  1837  	if err != nil {
  1838  		t.Fatalf("err: %v", err)
  1839  	}
  1840  
  1841  	jobs = gatherJobs(iter)
  1842  	if len(jobs) != 1 {
  1843  		t.Fatalf("err: %v", err)
  1844  	}
  1845  	if watchFired(ws) {
  1846  		t.Fatalf("bad")
  1847  	}
  1848  
  1849  	job = mock.Job()
  1850  	job.ID = "riak"
  1851  	err = state.UpsertJob(1001, job)
  1852  	if err != nil {
  1853  		t.Fatalf("err: %v", err)
  1854  	}
  1855  
  1856  	if !watchFired(ws) {
  1857  		t.Fatalf("bad")
  1858  	}
  1859  
  1860  	ws = memdb.NewWatchSet()
  1861  	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "r")
  1862  	if err != nil {
  1863  		t.Fatalf("err: %v", err)
  1864  	}
  1865  
  1866  	jobs = gatherJobs(iter)
  1867  	if len(jobs) != 2 {
  1868  		t.Fatalf("err: %v", err)
  1869  	}
  1870  
  1871  	iter, err = state.JobsByIDPrefix(ws, job.Namespace, "ri")
  1872  	if err != nil {
  1873  		t.Fatalf("err: %v", err)
  1874  	}
  1875  
  1876  	jobs = gatherJobs(iter)
  1877  	if len(jobs) != 1 {
  1878  		t.Fatalf("err: %v", err)
  1879  	}
  1880  	if watchFired(ws) {
  1881  		t.Fatalf("bad")
  1882  	}
  1883  }
  1884  
  1885  func TestStateStore_JobsByPeriodic(t *testing.T) {
  1886  	state := testStateStore(t)
  1887  	var periodic, nonPeriodic []*structs.Job
  1888  
  1889  	for i := 0; i < 10; i++ {
  1890  		job := mock.Job()
  1891  		nonPeriodic = append(nonPeriodic, job)
  1892  
  1893  		err := state.UpsertJob(1000+uint64(i), job)
  1894  		if err != nil {
  1895  			t.Fatalf("err: %v", err)
  1896  		}
  1897  	}
  1898  
  1899  	for i := 0; i < 10; i++ {
  1900  		job := mock.PeriodicJob()
  1901  		periodic = append(periodic, job)
  1902  
  1903  		err := state.UpsertJob(2000+uint64(i), job)
  1904  		if err != nil {
  1905  			t.Fatalf("err: %v", err)
  1906  		}
  1907  	}
  1908  
  1909  	ws := memdb.NewWatchSet()
  1910  	iter, err := state.JobsByPeriodic(ws, true)
  1911  	if err != nil {
  1912  		t.Fatalf("err: %v", err)
  1913  	}
  1914  
  1915  	var outPeriodic []*structs.Job
  1916  	for {
  1917  		raw := iter.Next()
  1918  		if raw == nil {
  1919  			break
  1920  		}
  1921  		outPeriodic = append(outPeriodic, raw.(*structs.Job))
  1922  	}
  1923  
  1924  	iter, err = state.JobsByPeriodic(ws, false)
  1925  	if err != nil {
  1926  		t.Fatalf("err: %v", err)
  1927  	}
  1928  
  1929  	var outNonPeriodic []*structs.Job
  1930  	for {
  1931  		raw := iter.Next()
  1932  		if raw == nil {
  1933  			break
  1934  		}
  1935  		outNonPeriodic = append(outNonPeriodic, raw.(*structs.Job))
  1936  	}
  1937  
  1938  	sort.Sort(JobIDSort(periodic))
  1939  	sort.Sort(JobIDSort(nonPeriodic))
  1940  	sort.Sort(JobIDSort(outPeriodic))
  1941  	sort.Sort(JobIDSort(outNonPeriodic))
  1942  
  1943  	if !reflect.DeepEqual(periodic, outPeriodic) {
  1944  		t.Fatalf("bad: %#v %#v", periodic, outPeriodic)
  1945  	}
  1946  
  1947  	if !reflect.DeepEqual(nonPeriodic, outNonPeriodic) {
  1948  		t.Fatalf("bad: %#v %#v", nonPeriodic, outNonPeriodic)
  1949  	}
  1950  	if watchFired(ws) {
  1951  		t.Fatalf("bad")
  1952  	}
  1953  }
  1954  
  1955  func TestStateStore_JobsByScheduler(t *testing.T) {
  1956  	state := testStateStore(t)
  1957  	var serviceJobs []*structs.Job
  1958  	var sysJobs []*structs.Job
  1959  
  1960  	for i := 0; i < 10; i++ {
  1961  		job := mock.Job()
  1962  		serviceJobs = append(serviceJobs, job)
  1963  
  1964  		err := state.UpsertJob(1000+uint64(i), job)
  1965  		if err != nil {
  1966  			t.Fatalf("err: %v", err)
  1967  		}
  1968  	}
  1969  
  1970  	for i := 0; i < 10; i++ {
  1971  		job := mock.SystemJob()
  1972  		job.Status = structs.JobStatusRunning
  1973  		sysJobs = append(sysJobs, job)
  1974  
  1975  		err := state.UpsertJob(2000+uint64(i), job)
  1976  		if err != nil {
  1977  			t.Fatalf("err: %v", err)
  1978  		}
  1979  	}
  1980  
  1981  	ws := memdb.NewWatchSet()
  1982  	iter, err := state.JobsByScheduler(ws, "service")
  1983  	if err != nil {
  1984  		t.Fatalf("err: %v", err)
  1985  	}
  1986  
  1987  	var outService []*structs.Job
  1988  	for {
  1989  		raw := iter.Next()
  1990  		if raw == nil {
  1991  			break
  1992  		}
  1993  		outService = append(outService, raw.(*structs.Job))
  1994  	}
  1995  
  1996  	iter, err = state.JobsByScheduler(ws, "system")
  1997  	if err != nil {
  1998  		t.Fatalf("err: %v", err)
  1999  	}
  2000  
  2001  	var outSystem []*structs.Job
  2002  	for {
  2003  		raw := iter.Next()
  2004  		if raw == nil {
  2005  			break
  2006  		}
  2007  		outSystem = append(outSystem, raw.(*structs.Job))
  2008  	}
  2009  
  2010  	sort.Sort(JobIDSort(serviceJobs))
  2011  	sort.Sort(JobIDSort(sysJobs))
  2012  	sort.Sort(JobIDSort(outService))
  2013  	sort.Sort(JobIDSort(outSystem))
  2014  
  2015  	if !reflect.DeepEqual(serviceJobs, outService) {
  2016  		t.Fatalf("bad: %#v %#v", serviceJobs, outService)
  2017  	}
  2018  
  2019  	if !reflect.DeepEqual(sysJobs, outSystem) {
  2020  		t.Fatalf("bad: %#v %#v", sysJobs, outSystem)
  2021  	}
  2022  	if watchFired(ws) {
  2023  		t.Fatalf("bad")
  2024  	}
  2025  }
  2026  
  2027  func TestStateStore_JobsByGC(t *testing.T) {
  2028  	state := testStateStore(t)
  2029  	gc, nonGc := make(map[string]struct{}), make(map[string]struct{})
  2030  
  2031  	for i := 0; i < 20; i++ {
  2032  		var job *structs.Job
  2033  		if i%2 == 0 {
  2034  			job = mock.Job()
  2035  		} else {
  2036  			job = mock.PeriodicJob()
  2037  		}
  2038  		nonGc[job.ID] = struct{}{}
  2039  
  2040  		if err := state.UpsertJob(1000+uint64(i), job); err != nil {
  2041  			t.Fatalf("err: %v", err)
  2042  		}
  2043  	}
  2044  
  2045  	for i := 0; i < 20; i += 2 {
  2046  		job := mock.Job()
  2047  		job.Type = structs.JobTypeBatch
  2048  		gc[job.ID] = struct{}{}
  2049  
  2050  		if err := state.UpsertJob(2000+uint64(i), job); err != nil {
  2051  			t.Fatalf("err: %v", err)
  2052  		}
  2053  
  2054  		// Create an eval for it
  2055  		eval := mock.Eval()
  2056  		eval.JobID = job.ID
  2057  		eval.Status = structs.EvalStatusComplete
  2058  		if err := state.UpsertEvals(2000+uint64(i+1), []*structs.Evaluation{eval}); err != nil {
  2059  			t.Fatalf("err: %v", err)
  2060  		}
  2061  
  2062  	}
  2063  
  2064  	ws := memdb.NewWatchSet()
  2065  	iter, err := state.JobsByGC(ws, true)
  2066  	if err != nil {
  2067  		t.Fatalf("err: %v", err)
  2068  	}
  2069  
  2070  	outGc := make(map[string]struct{})
  2071  	for i := iter.Next(); i != nil; i = iter.Next() {
  2072  		j := i.(*structs.Job)
  2073  		outGc[j.ID] = struct{}{}
  2074  	}
  2075  
  2076  	iter, err = state.JobsByGC(ws, false)
  2077  	if err != nil {
  2078  		t.Fatalf("err: %v", err)
  2079  	}
  2080  
  2081  	outNonGc := make(map[string]struct{})
  2082  	for i := iter.Next(); i != nil; i = iter.Next() {
  2083  		j := i.(*structs.Job)
  2084  		outNonGc[j.ID] = struct{}{}
  2085  	}
  2086  
  2087  	if !reflect.DeepEqual(gc, outGc) {
  2088  		t.Fatalf("bad: %#v %#v", gc, outGc)
  2089  	}
  2090  
  2091  	if !reflect.DeepEqual(nonGc, outNonGc) {
  2092  		t.Fatalf("bad: %#v %#v", nonGc, outNonGc)
  2093  	}
  2094  	if watchFired(ws) {
  2095  		t.Fatalf("bad")
  2096  	}
  2097  }
  2098  
  2099  func TestStateStore_RestoreJob(t *testing.T) {
  2100  	state := testStateStore(t)
  2101  	job := mock.Job()
  2102  
  2103  	restore, err := state.Restore()
  2104  	if err != nil {
  2105  		t.Fatalf("err: %v", err)
  2106  	}
  2107  
  2108  	err = restore.JobRestore(job)
  2109  	if err != nil {
  2110  		t.Fatalf("err: %v", err)
  2111  	}
  2112  	restore.Commit()
  2113  
  2114  	ws := memdb.NewWatchSet()
  2115  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  2116  	if err != nil {
  2117  		t.Fatalf("err: %v", err)
  2118  	}
  2119  
  2120  	if !reflect.DeepEqual(out, job) {
  2121  		t.Fatalf("Bad: %#v %#v", out, job)
  2122  	}
  2123  }
  2124  
  2125  // This test ensures that the state restore creates the EphemeralDisk for a job if
  2126  // it doesn't have one
  2127  // COMPAT 0.4.1 -> 0.5
  2128  func TestStateStore_Jobs_NoEphemeralDisk(t *testing.T) {
  2129  	state := testStateStore(t)
  2130  	job := mock.Job()
  2131  
  2132  	// Set EphemeralDisk to nil and set the DiskMB to 150
  2133  	job.TaskGroups[0].EphemeralDisk = nil
  2134  	job.TaskGroups[0].Tasks[0].Resources.DiskMB = 150
  2135  
  2136  	restore, err := state.Restore()
  2137  	if err != nil {
  2138  		t.Fatalf("err: %v", err)
  2139  	}
  2140  
  2141  	err = restore.JobRestore(job)
  2142  	if err != nil {
  2143  		t.Fatalf("err: %v", err)
  2144  	}
  2145  	restore.Commit()
  2146  
  2147  	ws := memdb.NewWatchSet()
  2148  	out, err := state.JobByID(ws, job.Namespace, job.ID)
  2149  	if err != nil {
  2150  		t.Fatalf("err: %v", err)
  2151  	}
  2152  
  2153  	// Expect job to have local disk and clear out the task's disk resource ask
  2154  	expected := job.Copy()
  2155  	expected.TaskGroups[0].EphemeralDisk = &structs.EphemeralDisk{
  2156  		SizeMB: 150,
  2157  	}
  2158  	expected.TaskGroups[0].Tasks[0].Resources.DiskMB = 0
  2159  	if !reflect.DeepEqual(out, expected) {
  2160  		t.Fatalf("Bad: %#v %#v", out, job)
  2161  	}
  2162  }
  2163  
  2164  func TestStateStore_UpsertPeriodicLaunch(t *testing.T) {
  2165  	state := testStateStore(t)
  2166  	job := mock.Job()
  2167  	launch := &structs.PeriodicLaunch{
  2168  		ID:        job.ID,
  2169  		Namespace: job.Namespace,
  2170  		Launch:    time.Now(),
  2171  	}
  2172  
  2173  	// Create a watchset so we can test that upsert fires the watch
  2174  	ws := memdb.NewWatchSet()
  2175  	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
  2176  		t.Fatalf("bad: %v", err)
  2177  	}
  2178  
  2179  	err := state.UpsertPeriodicLaunch(1000, launch)
  2180  	if err != nil {
  2181  		t.Fatalf("err: %v", err)
  2182  	}
  2183  
  2184  	if !watchFired(ws) {
  2185  		t.Fatalf("bad")
  2186  	}
  2187  
  2188  	ws = memdb.NewWatchSet()
  2189  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2190  	if err != nil {
  2191  		t.Fatalf("err: %v", err)
  2192  	}
  2193  	if out.CreateIndex != 1000 {
  2194  		t.Fatalf("bad: %#v", out)
  2195  	}
  2196  	if out.ModifyIndex != 1000 {
  2197  		t.Fatalf("bad: %#v", out)
  2198  	}
  2199  
  2200  	if !reflect.DeepEqual(launch, out) {
  2201  		t.Fatalf("bad: %#v %#v", job, out)
  2202  	}
  2203  
  2204  	index, err := state.Index("periodic_launch")
  2205  	if err != nil {
  2206  		t.Fatalf("err: %v", err)
  2207  	}
  2208  	if index != 1000 {
  2209  		t.Fatalf("bad: %d", index)
  2210  	}
  2211  
  2212  	if watchFired(ws) {
  2213  		t.Fatalf("bad")
  2214  	}
  2215  }
  2216  
  2217  func TestStateStore_UpdateUpsertPeriodicLaunch(t *testing.T) {
  2218  	state := testStateStore(t)
  2219  	job := mock.Job()
  2220  	launch := &structs.PeriodicLaunch{
  2221  		ID:        job.ID,
  2222  		Namespace: job.Namespace,
  2223  		Launch:    time.Now(),
  2224  	}
  2225  
  2226  	err := state.UpsertPeriodicLaunch(1000, launch)
  2227  	if err != nil {
  2228  		t.Fatalf("err: %v", err)
  2229  	}
  2230  
  2231  	// Create a watchset so we can test that upsert fires the watch
  2232  	ws := memdb.NewWatchSet()
  2233  	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
  2234  		t.Fatalf("bad: %v", err)
  2235  	}
  2236  
  2237  	launch2 := &structs.PeriodicLaunch{
  2238  		ID:        job.ID,
  2239  		Namespace: job.Namespace,
  2240  		Launch:    launch.Launch.Add(1 * time.Second),
  2241  	}
  2242  	err = state.UpsertPeriodicLaunch(1001, launch2)
  2243  	if err != nil {
  2244  		t.Fatalf("err: %v", err)
  2245  	}
  2246  
  2247  	if !watchFired(ws) {
  2248  		t.Fatalf("bad")
  2249  	}
  2250  
  2251  	ws = memdb.NewWatchSet()
  2252  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2253  	if err != nil {
  2254  		t.Fatalf("err: %v", err)
  2255  	}
  2256  	if out.CreateIndex != 1000 {
  2257  		t.Fatalf("bad: %#v", out)
  2258  	}
  2259  	if out.ModifyIndex != 1001 {
  2260  		t.Fatalf("bad: %#v", out)
  2261  	}
  2262  
  2263  	if !reflect.DeepEqual(launch2, out) {
  2264  		t.Fatalf("bad: %#v %#v", launch2, out)
  2265  	}
  2266  
  2267  	index, err := state.Index("periodic_launch")
  2268  	if err != nil {
  2269  		t.Fatalf("err: %v", err)
  2270  	}
  2271  	if index != 1001 {
  2272  		t.Fatalf("bad: %d", index)
  2273  	}
  2274  
  2275  	if watchFired(ws) {
  2276  		t.Fatalf("bad")
  2277  	}
  2278  }
  2279  
  2280  func TestStateStore_DeletePeriodicLaunch(t *testing.T) {
  2281  	state := testStateStore(t)
  2282  	job := mock.Job()
  2283  	launch := &structs.PeriodicLaunch{
  2284  		ID:        job.ID,
  2285  		Namespace: job.Namespace,
  2286  		Launch:    time.Now(),
  2287  	}
  2288  
  2289  	err := state.UpsertPeriodicLaunch(1000, launch)
  2290  	if err != nil {
  2291  		t.Fatalf("err: %v", err)
  2292  	}
  2293  
  2294  	// Create a watchset so we can test that delete fires the watch
  2295  	ws := memdb.NewWatchSet()
  2296  	if _, err := state.PeriodicLaunchByID(ws, job.Namespace, launch.ID); err != nil {
  2297  		t.Fatalf("bad: %v", err)
  2298  	}
  2299  
  2300  	err = state.DeletePeriodicLaunch(1001, launch.Namespace, launch.ID)
  2301  	if err != nil {
  2302  		t.Fatalf("err: %v", err)
  2303  	}
  2304  
  2305  	if !watchFired(ws) {
  2306  		t.Fatalf("bad")
  2307  	}
  2308  
  2309  	ws = memdb.NewWatchSet()
  2310  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2311  	if err != nil {
  2312  		t.Fatalf("err: %v", err)
  2313  	}
  2314  
  2315  	if out != nil {
  2316  		t.Fatalf("bad: %#v %#v", job, out)
  2317  	}
  2318  
  2319  	index, err := state.Index("periodic_launch")
  2320  	if err != nil {
  2321  		t.Fatalf("err: %v", err)
  2322  	}
  2323  	if index != 1001 {
  2324  		t.Fatalf("bad: %d", index)
  2325  	}
  2326  
  2327  	if watchFired(ws) {
  2328  		t.Fatalf("bad")
  2329  	}
  2330  }
  2331  
  2332  func TestStateStore_PeriodicLaunches(t *testing.T) {
  2333  	state := testStateStore(t)
  2334  	var launches []*structs.PeriodicLaunch
  2335  
  2336  	for i := 0; i < 10; i++ {
  2337  		job := mock.Job()
  2338  		launch := &structs.PeriodicLaunch{
  2339  			ID:        job.ID,
  2340  			Namespace: job.Namespace,
  2341  			Launch:    time.Now(),
  2342  		}
  2343  		launches = append(launches, launch)
  2344  
  2345  		err := state.UpsertPeriodicLaunch(1000+uint64(i), launch)
  2346  		if err != nil {
  2347  			t.Fatalf("err: %v", err)
  2348  		}
  2349  	}
  2350  
  2351  	ws := memdb.NewWatchSet()
  2352  	iter, err := state.PeriodicLaunches(ws)
  2353  	if err != nil {
  2354  		t.Fatalf("err: %v", err)
  2355  	}
  2356  
  2357  	out := make(map[string]*structs.PeriodicLaunch, 10)
  2358  	for {
  2359  		raw := iter.Next()
  2360  		if raw == nil {
  2361  			break
  2362  		}
  2363  		launch := raw.(*structs.PeriodicLaunch)
  2364  		if _, ok := out[launch.ID]; ok {
  2365  			t.Fatalf("duplicate: %v", launch.ID)
  2366  		}
  2367  
  2368  		out[launch.ID] = launch
  2369  	}
  2370  
  2371  	for _, launch := range launches {
  2372  		l, ok := out[launch.ID]
  2373  		if !ok {
  2374  			t.Fatalf("bad %v", launch.ID)
  2375  		}
  2376  
  2377  		if !reflect.DeepEqual(launch, l) {
  2378  			t.Fatalf("bad: %#v %#v", launch, l)
  2379  		}
  2380  
  2381  		delete(out, launch.ID)
  2382  	}
  2383  
  2384  	if len(out) != 0 {
  2385  		t.Fatalf("leftover: %#v", out)
  2386  	}
  2387  
  2388  	if watchFired(ws) {
  2389  		t.Fatalf("bad")
  2390  	}
  2391  }
  2392  
  2393  func TestStateStore_RestorePeriodicLaunch(t *testing.T) {
  2394  	state := testStateStore(t)
  2395  	job := mock.Job()
  2396  	launch := &structs.PeriodicLaunch{
  2397  		ID:        job.ID,
  2398  		Namespace: job.Namespace,
  2399  		Launch:    time.Now(),
  2400  	}
  2401  
  2402  	restore, err := state.Restore()
  2403  	if err != nil {
  2404  		t.Fatalf("err: %v", err)
  2405  	}
  2406  
  2407  	err = restore.PeriodicLaunchRestore(launch)
  2408  	if err != nil {
  2409  		t.Fatalf("err: %v", err)
  2410  	}
  2411  	restore.Commit()
  2412  
  2413  	ws := memdb.NewWatchSet()
  2414  	out, err := state.PeriodicLaunchByID(ws, job.Namespace, job.ID)
  2415  	if err != nil {
  2416  		t.Fatalf("err: %v", err)
  2417  	}
  2418  
  2419  	if !reflect.DeepEqual(out, launch) {
  2420  		t.Fatalf("Bad: %#v %#v", out, job)
  2421  	}
  2422  
  2423  	if watchFired(ws) {
  2424  		t.Fatalf("bad")
  2425  	}
  2426  }
  2427  
  2428  func TestStateStore_RestoreJobVersion(t *testing.T) {
  2429  	state := testStateStore(t)
  2430  	job := mock.Job()
  2431  
  2432  	restore, err := state.Restore()
  2433  	if err != nil {
  2434  		t.Fatalf("err: %v", err)
  2435  	}
  2436  
  2437  	err = restore.JobVersionRestore(job)
  2438  	if err != nil {
  2439  		t.Fatalf("err: %v", err)
  2440  	}
  2441  	restore.Commit()
  2442  
  2443  	ws := memdb.NewWatchSet()
  2444  	out, err := state.JobByIDAndVersion(ws, job.Namespace, job.ID, job.Version)
  2445  	if err != nil {
  2446  		t.Fatalf("err: %v", err)
  2447  	}
  2448  
  2449  	if !reflect.DeepEqual(out, job) {
  2450  		t.Fatalf("Bad: %#v %#v", out, job)
  2451  	}
  2452  
  2453  	if watchFired(ws) {
  2454  		t.Fatalf("bad")
  2455  	}
  2456  }
  2457  
  2458  func TestStateStore_RestoreDeployment(t *testing.T) {
  2459  	state := testStateStore(t)
  2460  	d := mock.Deployment()
  2461  
  2462  	restore, err := state.Restore()
  2463  	if err != nil {
  2464  		t.Fatalf("err: %v", err)
  2465  	}
  2466  
  2467  	err = restore.DeploymentRestore(d)
  2468  	if err != nil {
  2469  		t.Fatalf("err: %v", err)
  2470  	}
  2471  	restore.Commit()
  2472  
  2473  	ws := memdb.NewWatchSet()
  2474  	out, err := state.DeploymentByID(ws, d.ID)
  2475  	if err != nil {
  2476  		t.Fatalf("err: %v", err)
  2477  	}
  2478  
  2479  	if !reflect.DeepEqual(out, d) {
  2480  		t.Fatalf("Bad: %#v %#v", out, d)
  2481  	}
  2482  
  2483  	if watchFired(ws) {
  2484  		t.Fatalf("bad")
  2485  	}
  2486  }
  2487  
  2488  func TestStateStore_RestoreJobSummary(t *testing.T) {
  2489  	state := testStateStore(t)
  2490  	job := mock.Job()
  2491  	jobSummary := &structs.JobSummary{
  2492  		JobID:     job.ID,
  2493  		Namespace: job.Namespace,
  2494  		Summary: map[string]structs.TaskGroupSummary{
  2495  			"web": {
  2496  				Starting: 10,
  2497  			},
  2498  		},
  2499  	}
  2500  	restore, err := state.Restore()
  2501  	if err != nil {
  2502  		t.Fatalf("err: %v", err)
  2503  	}
  2504  
  2505  	err = restore.JobSummaryRestore(jobSummary)
  2506  	if err != nil {
  2507  		t.Fatalf("err: %v", err)
  2508  	}
  2509  	restore.Commit()
  2510  
  2511  	ws := memdb.NewWatchSet()
  2512  	out, err := state.JobSummaryByID(ws, job.Namespace, job.ID)
  2513  	if err != nil {
  2514  		t.Fatalf("err: %v", err)
  2515  	}
  2516  
  2517  	if !reflect.DeepEqual(out, jobSummary) {
  2518  		t.Fatalf("Bad: %#v %#v", out, jobSummary)
  2519  	}
  2520  }
  2521  
  2522  func TestStateStore_Indexes(t *testing.T) {
  2523  	state := testStateStore(t)
  2524  	node := mock.Node()
  2525  
  2526  	err := state.UpsertNode(1000, node)
  2527  	if err != nil {
  2528  		t.Fatalf("err: %v", err)
  2529  	}
  2530  
  2531  	iter, err := state.Indexes()
  2532  	if err != nil {
  2533  		t.Fatalf("err: %v", err)
  2534  	}
  2535  
  2536  	var out []*IndexEntry
  2537  	for {
  2538  		raw := iter.Next()
  2539  		if raw == nil {
  2540  			break
  2541  		}
  2542  		out = append(out, raw.(*IndexEntry))
  2543  	}
  2544  
  2545  	expect := &IndexEntry{"nodes", 1000}
  2546  	if l := len(out); l != 1 && l != 2 {
  2547  		t.Fatalf("unexpected number of index entries: %v", out)
  2548  	}
  2549  
  2550  	for _, index := range out {
  2551  		if index.Key != expect.Key {
  2552  			continue
  2553  		}
  2554  		if index.Value != expect.Value {
  2555  			t.Fatalf("bad index; got %d; want %d", index.Value, expect.Value)
  2556  		}
  2557  
  2558  		// We matched
  2559  		return
  2560  	}
  2561  
  2562  	t.Fatal("did not find expected index entry")
  2563  }
  2564  
  2565  func TestStateStore_LatestIndex(t *testing.T) {
  2566  	state := testStateStore(t)
  2567  
  2568  	if err := state.UpsertNode(1000, mock.Node()); err != nil {
  2569  		t.Fatalf("err: %v", err)
  2570  	}
  2571  
  2572  	exp := uint64(2000)
  2573  	if err := state.UpsertJob(exp, mock.Job()); err != nil {
  2574  		t.Fatalf("err: %v", err)
  2575  	}
  2576  
  2577  	latest, err := state.LatestIndex()
  2578  	if err != nil {
  2579  		t.Fatalf("err: %v", err)
  2580  	}
  2581  
  2582  	if latest != exp {
  2583  		t.Fatalf("LatestIndex() returned %d; want %d", latest, exp)
  2584  	}
  2585  }
  2586  
  2587  func TestStateStore_RestoreIndex(t *testing.T) {
  2588  	state := testStateStore(t)
  2589  
  2590  	restore, err := state.Restore()
  2591  	if err != nil {
  2592  		t.Fatalf("err: %v", err)
  2593  	}
  2594  
  2595  	index := &IndexEntry{"jobs", 1000}
  2596  	err = restore.IndexRestore(index)
  2597  	if err != nil {
  2598  		t.Fatalf("err: %v", err)
  2599  	}
  2600  
  2601  	restore.Commit()
  2602  
  2603  	out, err := state.Index("jobs")
  2604  	if err != nil {
  2605  		t.Fatalf("err: %v", err)
  2606  	}
  2607  
  2608  	if out != 1000 {
  2609  		t.Fatalf("Bad: %#v %#v", out, 1000)
  2610  	}
  2611  }
  2612  
  2613  func TestStateStore_UpsertEvals_Eval(t *testing.T) {
  2614  	state := testStateStore(t)
  2615  	eval := mock.Eval()
  2616  
  2617  	// Create a watchset so we can test that upsert fires the watch
  2618  	ws := memdb.NewWatchSet()
  2619  	if _, err := state.EvalByID(ws, eval.ID); err != nil {
  2620  		t.Fatalf("bad: %v", err)
  2621  	}
  2622  
  2623  	err := state.UpsertEvals(1000, []*structs.Evaluation{eval})
  2624  	if err != nil {
  2625  		t.Fatalf("err: %v", err)
  2626  	}
  2627  
  2628  	if !watchFired(ws) {
  2629  		t.Fatalf("bad")
  2630  	}
  2631  
  2632  	ws = memdb.NewWatchSet()
  2633  	out, err := state.EvalByID(ws, eval.ID)
  2634  	if err != nil {
  2635  		t.Fatalf("err: %v", err)
  2636  	}
  2637  
  2638  	if !reflect.DeepEqual(eval, out) {
  2639  		t.Fatalf("bad: %#v %#v", eval, out)
  2640  	}
  2641  
  2642  	index, err := state.Index("evals")
  2643  	if err != nil {
  2644  		t.Fatalf("err: %v", err)
  2645  	}
  2646  	if index != 1000 {
  2647  		t.Fatalf("bad: %d", index)
  2648  	}
  2649  
  2650  	if watchFired(ws) {
  2651  		t.Fatalf("bad")
  2652  	}
  2653  }
  2654  
  2655  func TestStateStore_UpsertEvals_CancelBlocked(t *testing.T) {
  2656  	state := testStateStore(t)
  2657  
  2658  	// Create two blocked evals for the same job
  2659  	j := "test-job"
  2660  	b1, b2 := mock.Eval(), mock.Eval()
  2661  	b1.JobID = j
  2662  	b1.Status = structs.EvalStatusBlocked
  2663  	b2.JobID = j
  2664  	b2.Status = structs.EvalStatusBlocked
  2665  
  2666  	err := state.UpsertEvals(999, []*structs.Evaluation{b1, b2})
  2667  	if err != nil {
  2668  		t.Fatalf("err: %v", err)
  2669  	}
  2670  
  2671  	// Create one complete and successful eval for the job
  2672  	eval := mock.Eval()
  2673  	eval.JobID = j
  2674  	eval.Status = structs.EvalStatusComplete
  2675  
  2676  	// Create a watchset so we can test that the upsert of the complete eval
  2677  	// fires the watch
  2678  	ws := memdb.NewWatchSet()
  2679  	if _, err := state.EvalByID(ws, b1.ID); err != nil {
  2680  		t.Fatalf("bad: %v", err)
  2681  	}
  2682  	if _, err := state.EvalByID(ws, b2.ID); err != nil {
  2683  		t.Fatalf("bad: %v", err)
  2684  	}
  2685  
  2686  	if err := state.UpsertEvals(1000, []*structs.Evaluation{eval}); err != nil {
  2687  		t.Fatalf("err: %v", err)
  2688  	}
  2689  
  2690  	if !watchFired(ws) {
  2691  		t.Fatalf("bad")
  2692  	}
  2693  
  2694  	ws = memdb.NewWatchSet()
  2695  	out, err := state.EvalByID(ws, eval.ID)
  2696  	if err != nil {
  2697  		t.Fatalf("err: %v", err)
  2698  	}
  2699  
  2700  	if !reflect.DeepEqual(eval, out) {
  2701  		t.Fatalf("bad: %#v %#v", eval, out)
  2702  	}
  2703  
  2704  	index, err := state.Index("evals")
  2705  	if err != nil {
  2706  		t.Fatalf("err: %v", err)
  2707  	}
  2708  	if index != 1000 {
  2709  		t.Fatalf("bad: %d", index)
  2710  	}
  2711  
  2712  	// Get b1/b2 and check they are cancelled
  2713  	out1, err := state.EvalByID(ws, b1.ID)
  2714  	if err != nil {
  2715  		t.Fatalf("err: %v", err)
  2716  	}
  2717  
  2718  	out2, err := state.EvalByID(ws, b2.ID)
  2719  	if err != nil {
  2720  		t.Fatalf("err: %v", err)
  2721  	}
  2722  
  2723  	if out1.Status != structs.EvalStatusCancelled || out2.Status != structs.EvalStatusCancelled {
  2724  		t.Fatalf("bad: %#v %#v", out1, out2)
  2725  	}
  2726  
  2727  	if watchFired(ws) {
  2728  		t.Fatalf("bad")
  2729  	}
  2730  }
  2731  
  2732  func TestStateStore_Update_UpsertEvals_Eval(t *testing.T) {
  2733  	state := testStateStore(t)
  2734  	eval := mock.Eval()
  2735  
  2736  	err := state.UpsertEvals(1000, []*structs.Evaluation{eval})
  2737  	if err != nil {
  2738  		t.Fatalf("err: %v", err)
  2739  	}
  2740  
  2741  	// Create a watchset so we can test that delete fires the watch
  2742  	ws := memdb.NewWatchSet()
  2743  	ws2 := memdb.NewWatchSet()
  2744  	if _, err := state.EvalByID(ws, eval.ID); err != nil {
  2745  		t.Fatalf("bad: %v", err)
  2746  	}
  2747  
  2748  	if _, err := state.EvalsByJob(ws2, eval.Namespace, eval.JobID); err != nil {
  2749  		t.Fatalf("bad: %v", err)
  2750  	}
  2751  
  2752  	eval2 := mock.Eval()
  2753  	eval2.ID = eval.ID
  2754  	eval2.JobID = eval.JobID
  2755  	err = state.UpsertEvals(1001, []*structs.Evaluation{eval2})
  2756  	if err != nil {
  2757  		t.Fatalf("err: %v", err)
  2758  	}
  2759  
  2760  	if !watchFired(ws) {
  2761  		t.Fatalf("bad")
  2762  	}
  2763  	if !watchFired(ws2) {
  2764  		t.Fatalf("bad")
  2765  	}
  2766  
  2767  	ws = memdb.NewWatchSet()
  2768  	out, err := state.EvalByID(ws, eval.ID)
  2769  	if err != nil {
  2770  		t.Fatalf("err: %v", err)
  2771  	}
  2772  
  2773  	if !reflect.DeepEqual(eval2, out) {
  2774  		t.Fatalf("bad: %#v %#v", eval2, out)
  2775  	}
  2776  
  2777  	if out.CreateIndex != 1000 {
  2778  		t.Fatalf("bad: %#v", out)
  2779  	}
  2780  	if out.ModifyIndex != 1001 {
  2781  		t.Fatalf("bad: %#v", out)
  2782  	}
  2783  
  2784  	index, err := state.Index("evals")
  2785  	if err != nil {
  2786  		t.Fatalf("err: %v", err)
  2787  	}
  2788  	if index != 1001 {
  2789  		t.Fatalf("bad: %d", index)
  2790  	}
  2791  
  2792  	if watchFired(ws) {
  2793  		t.Fatalf("bad")
  2794  	}
  2795  }
  2796  
  2797  func TestStateStore_UpsertEvals_Eval_ChildJob(t *testing.T) {
  2798  	state := testStateStore(t)
  2799  
  2800  	parent := mock.Job()
  2801  	if err := state.UpsertJob(998, parent); err != nil {
  2802  		t.Fatalf("err: %v", err)
  2803  	}
  2804  
  2805  	child := mock.Job()
  2806  	child.ParentID = parent.ID
  2807  
  2808  	if err := state.UpsertJob(999, child); err != nil {
  2809  		t.Fatalf("err: %v", err)
  2810  	}
  2811  
  2812  	eval := mock.Eval()
  2813  	eval.Status = structs.EvalStatusComplete
  2814  	eval.JobID = child.ID
  2815  
  2816  	// Create watchsets so we can test that upsert fires the watch
  2817  	ws := memdb.NewWatchSet()
  2818  	ws2 := memdb.NewWatchSet()
  2819  	ws3 := memdb.NewWatchSet()
  2820  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  2821  		t.Fatalf("bad: %v", err)
  2822  	}
  2823  	if _, err := state.EvalByID(ws2, eval.ID); err != nil {
  2824  		t.Fatalf("bad: %v", err)
  2825  	}
  2826  	if _, err := state.EvalsByJob(ws3, eval.Namespace, eval.JobID); err != nil {
  2827  		t.Fatalf("bad: %v", err)
  2828  	}
  2829  
  2830  	err := state.UpsertEvals(1000, []*structs.Evaluation{eval})
  2831  	if err != nil {
  2832  		t.Fatalf("err: %v", err)
  2833  	}
  2834  
  2835  	if !watchFired(ws) {
  2836  		t.Fatalf("bad")
  2837  	}
  2838  	if !watchFired(ws2) {
  2839  		t.Fatalf("bad")
  2840  	}
  2841  	if !watchFired(ws3) {
  2842  		t.Fatalf("bad")
  2843  	}
  2844  
  2845  	ws = memdb.NewWatchSet()
  2846  	out, err := state.EvalByID(ws, eval.ID)
  2847  	if err != nil {
  2848  		t.Fatalf("err: %v", err)
  2849  	}
  2850  
  2851  	if !reflect.DeepEqual(eval, out) {
  2852  		t.Fatalf("bad: %#v %#v", eval, out)
  2853  	}
  2854  
  2855  	index, err := state.Index("evals")
  2856  	if err != nil {
  2857  		t.Fatalf("err: %v", err)
  2858  	}
  2859  	if index != 1000 {
  2860  		t.Fatalf("bad: %d", index)
  2861  	}
  2862  
  2863  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  2864  	if err != nil {
  2865  		t.Fatalf("err: %v", err)
  2866  	}
  2867  	if summary == nil {
  2868  		t.Fatalf("nil summary")
  2869  	}
  2870  	if summary.JobID != parent.ID {
  2871  		t.Fatalf("bad summary id: %v", parent.ID)
  2872  	}
  2873  	if summary.Children == nil {
  2874  		t.Fatalf("nil children summary")
  2875  	}
  2876  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  2877  		t.Fatalf("bad children summary: %v", summary.Children)
  2878  	}
  2879  
  2880  	if watchFired(ws) {
  2881  		t.Fatalf("bad")
  2882  	}
  2883  }
  2884  
  2885  func TestStateStore_DeleteEval_Eval(t *testing.T) {
  2886  	state := testStateStore(t)
  2887  	eval1 := mock.Eval()
  2888  	eval2 := mock.Eval()
  2889  	alloc1 := mock.Alloc()
  2890  	alloc2 := mock.Alloc()
  2891  
  2892  	// Create watchsets so we can test that upsert fires the watch
  2893  	watches := make([]memdb.WatchSet, 12)
  2894  	for i := 0; i < 12; i++ {
  2895  		watches[i] = memdb.NewWatchSet()
  2896  	}
  2897  	if _, err := state.EvalByID(watches[0], eval1.ID); err != nil {
  2898  		t.Fatalf("bad: %v", err)
  2899  	}
  2900  	if _, err := state.EvalByID(watches[1], eval2.ID); err != nil {
  2901  		t.Fatalf("bad: %v", err)
  2902  	}
  2903  	if _, err := state.EvalsByJob(watches[2], eval1.Namespace, eval1.JobID); err != nil {
  2904  		t.Fatalf("bad: %v", err)
  2905  	}
  2906  	if _, err := state.EvalsByJob(watches[3], eval2.Namespace, eval2.JobID); err != nil {
  2907  		t.Fatalf("bad: %v", err)
  2908  	}
  2909  	if _, err := state.AllocByID(watches[4], alloc1.ID); err != nil {
  2910  		t.Fatalf("bad: %v", err)
  2911  	}
  2912  	if _, err := state.AllocByID(watches[5], alloc2.ID); err != nil {
  2913  		t.Fatalf("bad: %v", err)
  2914  	}
  2915  	if _, err := state.AllocsByEval(watches[6], alloc1.EvalID); err != nil {
  2916  		t.Fatalf("bad: %v", err)
  2917  	}
  2918  	if _, err := state.AllocsByEval(watches[7], alloc2.EvalID); err != nil {
  2919  		t.Fatalf("bad: %v", err)
  2920  	}
  2921  	if _, err := state.AllocsByJob(watches[8], alloc1.Namespace, alloc1.JobID, false); err != nil {
  2922  		t.Fatalf("bad: %v", err)
  2923  	}
  2924  	if _, err := state.AllocsByJob(watches[9], alloc2.Namespace, alloc2.JobID, false); err != nil {
  2925  		t.Fatalf("bad: %v", err)
  2926  	}
  2927  	if _, err := state.AllocsByNode(watches[10], alloc1.NodeID); err != nil {
  2928  		t.Fatalf("bad: %v", err)
  2929  	}
  2930  	if _, err := state.AllocsByNode(watches[11], alloc2.NodeID); err != nil {
  2931  		t.Fatalf("bad: %v", err)
  2932  	}
  2933  
  2934  	state.UpsertJobSummary(900, mock.JobSummary(eval1.JobID))
  2935  	state.UpsertJobSummary(901, mock.JobSummary(eval2.JobID))
  2936  	state.UpsertJobSummary(902, mock.JobSummary(alloc1.JobID))
  2937  	state.UpsertJobSummary(903, mock.JobSummary(alloc2.JobID))
  2938  	err := state.UpsertEvals(1000, []*structs.Evaluation{eval1, eval2})
  2939  	if err != nil {
  2940  		t.Fatalf("err: %v", err)
  2941  	}
  2942  
  2943  	err = state.UpsertAllocs(1001, []*structs.Allocation{alloc1, alloc2})
  2944  	if err != nil {
  2945  		t.Fatalf("err: %v", err)
  2946  	}
  2947  
  2948  	err = state.DeleteEval(1002, []string{eval1.ID, eval2.ID}, []string{alloc1.ID, alloc2.ID})
  2949  	if err != nil {
  2950  		t.Fatalf("err: %v", err)
  2951  	}
  2952  
  2953  	for i, ws := range watches {
  2954  		if !watchFired(ws) {
  2955  			t.Fatalf("bad %d", i)
  2956  		}
  2957  	}
  2958  
  2959  	ws := memdb.NewWatchSet()
  2960  	out, err := state.EvalByID(ws, eval1.ID)
  2961  	if err != nil {
  2962  		t.Fatalf("err: %v", err)
  2963  	}
  2964  
  2965  	if out != nil {
  2966  		t.Fatalf("bad: %#v %#v", eval1, out)
  2967  	}
  2968  
  2969  	out, err = state.EvalByID(ws, eval2.ID)
  2970  	if err != nil {
  2971  		t.Fatalf("err: %v", err)
  2972  	}
  2973  
  2974  	if out != nil {
  2975  		t.Fatalf("bad: %#v %#v", eval1, out)
  2976  	}
  2977  
  2978  	outA, err := state.AllocByID(ws, alloc1.ID)
  2979  	if err != nil {
  2980  		t.Fatalf("err: %v", err)
  2981  	}
  2982  
  2983  	if out != nil {
  2984  		t.Fatalf("bad: %#v %#v", alloc1, outA)
  2985  	}
  2986  
  2987  	outA, err = state.AllocByID(ws, alloc2.ID)
  2988  	if err != nil {
  2989  		t.Fatalf("err: %v", err)
  2990  	}
  2991  
  2992  	if out != nil {
  2993  		t.Fatalf("bad: %#v %#v", alloc1, outA)
  2994  	}
  2995  
  2996  	index, err := state.Index("evals")
  2997  	if err != nil {
  2998  		t.Fatalf("err: %v", err)
  2999  	}
  3000  	if index != 1002 {
  3001  		t.Fatalf("bad: %d", index)
  3002  	}
  3003  
  3004  	index, err = state.Index("allocs")
  3005  	if err != nil {
  3006  		t.Fatalf("err: %v", err)
  3007  	}
  3008  	if index != 1002 {
  3009  		t.Fatalf("bad: %d", index)
  3010  	}
  3011  
  3012  	if watchFired(ws) {
  3013  		t.Fatalf("bad")
  3014  	}
  3015  }
  3016  
  3017  func TestStateStore_DeleteEval_ChildJob(t *testing.T) {
  3018  	state := testStateStore(t)
  3019  
  3020  	parent := mock.Job()
  3021  	if err := state.UpsertJob(998, parent); err != nil {
  3022  		t.Fatalf("err: %v", err)
  3023  	}
  3024  
  3025  	child := mock.Job()
  3026  	child.ParentID = parent.ID
  3027  
  3028  	if err := state.UpsertJob(999, child); err != nil {
  3029  		t.Fatalf("err: %v", err)
  3030  	}
  3031  
  3032  	eval1 := mock.Eval()
  3033  	eval1.JobID = child.ID
  3034  	alloc1 := mock.Alloc()
  3035  	alloc1.JobID = child.ID
  3036  
  3037  	err := state.UpsertEvals(1000, []*structs.Evaluation{eval1})
  3038  	if err != nil {
  3039  		t.Fatalf("err: %v", err)
  3040  	}
  3041  
  3042  	err = state.UpsertAllocs(1001, []*structs.Allocation{alloc1})
  3043  	if err != nil {
  3044  		t.Fatalf("err: %v", err)
  3045  	}
  3046  
  3047  	// Create watchsets so we can test that delete fires the watch
  3048  	ws := memdb.NewWatchSet()
  3049  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  3050  		t.Fatalf("bad: %v", err)
  3051  	}
  3052  
  3053  	err = state.DeleteEval(1002, []string{eval1.ID}, []string{alloc1.ID})
  3054  	if err != nil {
  3055  		t.Fatalf("err: %v", err)
  3056  	}
  3057  
  3058  	if !watchFired(ws) {
  3059  		t.Fatalf("bad")
  3060  	}
  3061  
  3062  	ws = memdb.NewWatchSet()
  3063  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  3064  	if err != nil {
  3065  		t.Fatalf("err: %v", err)
  3066  	}
  3067  	if summary == nil {
  3068  		t.Fatalf("nil summary")
  3069  	}
  3070  	if summary.JobID != parent.ID {
  3071  		t.Fatalf("bad summary id: %v", parent.ID)
  3072  	}
  3073  	if summary.Children == nil {
  3074  		t.Fatalf("nil children summary")
  3075  	}
  3076  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  3077  		t.Fatalf("bad children summary: %v", summary.Children)
  3078  	}
  3079  
  3080  	if watchFired(ws) {
  3081  		t.Fatalf("bad")
  3082  	}
  3083  }
  3084  
  3085  func TestStateStore_EvalsByJob(t *testing.T) {
  3086  	state := testStateStore(t)
  3087  
  3088  	eval1 := mock.Eval()
  3089  	eval2 := mock.Eval()
  3090  	eval2.JobID = eval1.JobID
  3091  	eval3 := mock.Eval()
  3092  	evals := []*structs.Evaluation{eval1, eval2}
  3093  
  3094  	err := state.UpsertEvals(1000, evals)
  3095  	if err != nil {
  3096  		t.Fatalf("err: %v", err)
  3097  	}
  3098  	err = state.UpsertEvals(1001, []*structs.Evaluation{eval3})
  3099  	if err != nil {
  3100  		t.Fatalf("err: %v", err)
  3101  	}
  3102  
  3103  	ws := memdb.NewWatchSet()
  3104  	out, err := state.EvalsByJob(ws, eval1.Namespace, eval1.JobID)
  3105  	if err != nil {
  3106  		t.Fatalf("err: %v", err)
  3107  	}
  3108  
  3109  	sort.Sort(EvalIDSort(evals))
  3110  	sort.Sort(EvalIDSort(out))
  3111  
  3112  	if !reflect.DeepEqual(evals, out) {
  3113  		t.Fatalf("bad: %#v %#v", evals, out)
  3114  	}
  3115  
  3116  	if watchFired(ws) {
  3117  		t.Fatalf("bad")
  3118  	}
  3119  }
  3120  
  3121  func TestStateStore_Evals(t *testing.T) {
  3122  	state := testStateStore(t)
  3123  	var evals []*structs.Evaluation
  3124  
  3125  	for i := 0; i < 10; i++ {
  3126  		eval := mock.Eval()
  3127  		evals = append(evals, eval)
  3128  
  3129  		err := state.UpsertEvals(1000+uint64(i), []*structs.Evaluation{eval})
  3130  		if err != nil {
  3131  			t.Fatalf("err: %v", err)
  3132  		}
  3133  	}
  3134  
  3135  	ws := memdb.NewWatchSet()
  3136  	iter, err := state.Evals(ws)
  3137  	if err != nil {
  3138  		t.Fatalf("err: %v", err)
  3139  	}
  3140  
  3141  	var out []*structs.Evaluation
  3142  	for {
  3143  		raw := iter.Next()
  3144  		if raw == nil {
  3145  			break
  3146  		}
  3147  		out = append(out, raw.(*structs.Evaluation))
  3148  	}
  3149  
  3150  	sort.Sort(EvalIDSort(evals))
  3151  	sort.Sort(EvalIDSort(out))
  3152  
  3153  	if !reflect.DeepEqual(evals, out) {
  3154  		t.Fatalf("bad: %#v %#v", evals, out)
  3155  	}
  3156  
  3157  	if watchFired(ws) {
  3158  		t.Fatalf("bad")
  3159  	}
  3160  }
  3161  
  3162  func TestStateStore_EvalsByIDPrefix(t *testing.T) {
  3163  	state := testStateStore(t)
  3164  	var evals []*structs.Evaluation
  3165  
  3166  	ids := []string{
  3167  		"aaaaaaaa-7bfb-395d-eb95-0685af2176b2",
  3168  		"aaaaaaab-7bfb-395d-eb95-0685af2176b2",
  3169  		"aaaaaabb-7bfb-395d-eb95-0685af2176b2",
  3170  		"aaaaabbb-7bfb-395d-eb95-0685af2176b2",
  3171  		"aaaabbbb-7bfb-395d-eb95-0685af2176b2",
  3172  		"aaabbbbb-7bfb-395d-eb95-0685af2176b2",
  3173  		"aabbbbbb-7bfb-395d-eb95-0685af2176b2",
  3174  		"abbbbbbb-7bfb-395d-eb95-0685af2176b2",
  3175  		"bbbbbbbb-7bfb-395d-eb95-0685af2176b2",
  3176  	}
  3177  	for i := 0; i < 9; i++ {
  3178  		eval := mock.Eval()
  3179  		eval.ID = ids[i]
  3180  		evals = append(evals, eval)
  3181  	}
  3182  
  3183  	err := state.UpsertEvals(1000, evals)
  3184  	if err != nil {
  3185  		t.Fatalf("err: %v", err)
  3186  	}
  3187  
  3188  	ws := memdb.NewWatchSet()
  3189  	iter, err := state.EvalsByIDPrefix(ws, structs.DefaultNamespace, "aaaa")
  3190  	if err != nil {
  3191  		t.Fatalf("err: %v", err)
  3192  	}
  3193  
  3194  	gatherEvals := func(iter memdb.ResultIterator) []*structs.Evaluation {
  3195  		var evals []*structs.Evaluation
  3196  		for {
  3197  			raw := iter.Next()
  3198  			if raw == nil {
  3199  				break
  3200  			}
  3201  			evals = append(evals, raw.(*structs.Evaluation))
  3202  		}
  3203  		return evals
  3204  	}
  3205  
  3206  	out := gatherEvals(iter)
  3207  	if len(out) != 5 {
  3208  		t.Fatalf("bad: expected five evaluations, got: %#v", out)
  3209  	}
  3210  
  3211  	sort.Sort(EvalIDSort(evals))
  3212  
  3213  	for index, eval := range out {
  3214  		if ids[index] != eval.ID {
  3215  			t.Fatalf("bad: got unexpected id: %s", eval.ID)
  3216  		}
  3217  	}
  3218  
  3219  	iter, err = state.EvalsByIDPrefix(ws, structs.DefaultNamespace, "b-a7bfb")
  3220  	if err != nil {
  3221  		t.Fatalf("err: %v", err)
  3222  	}
  3223  
  3224  	out = gatherEvals(iter)
  3225  	if len(out) != 0 {
  3226  		t.Fatalf("bad: unexpected zero evaluations, got: %#v", out)
  3227  	}
  3228  
  3229  	if watchFired(ws) {
  3230  		t.Fatalf("bad")
  3231  	}
  3232  }
  3233  
  3234  func TestStateStore_RestoreEval(t *testing.T) {
  3235  	state := testStateStore(t)
  3236  	eval := mock.Eval()
  3237  
  3238  	restore, err := state.Restore()
  3239  	if err != nil {
  3240  		t.Fatalf("err: %v", err)
  3241  	}
  3242  
  3243  	err = restore.EvalRestore(eval)
  3244  	if err != nil {
  3245  		t.Fatalf("err: %v", err)
  3246  	}
  3247  	restore.Commit()
  3248  
  3249  	ws := memdb.NewWatchSet()
  3250  	out, err := state.EvalByID(ws, eval.ID)
  3251  	if err != nil {
  3252  		t.Fatalf("err: %v", err)
  3253  	}
  3254  
  3255  	if !reflect.DeepEqual(out, eval) {
  3256  		t.Fatalf("Bad: %#v %#v", out, eval)
  3257  	}
  3258  }
  3259  
  3260  func TestStateStore_UpdateAllocsFromClient(t *testing.T) {
  3261  	state := testStateStore(t)
  3262  	parent := mock.Job()
  3263  	if err := state.UpsertJob(998, parent); err != nil {
  3264  		t.Fatalf("err: %v", err)
  3265  	}
  3266  
  3267  	child := mock.Job()
  3268  	child.ParentID = parent.ID
  3269  	if err := state.UpsertJob(999, child); err != nil {
  3270  		t.Fatalf("err: %v", err)
  3271  	}
  3272  
  3273  	alloc := mock.Alloc()
  3274  	alloc.JobID = child.ID
  3275  	alloc.Job = child
  3276  
  3277  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3278  	if err != nil {
  3279  		t.Fatalf("err: %v", err)
  3280  	}
  3281  
  3282  	ws := memdb.NewWatchSet()
  3283  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  3284  	if err != nil {
  3285  		t.Fatalf("err: %v", err)
  3286  	}
  3287  	if summary == nil {
  3288  		t.Fatalf("nil summary")
  3289  	}
  3290  	if summary.JobID != parent.ID {
  3291  		t.Fatalf("bad summary id: %v", parent.ID)
  3292  	}
  3293  	if summary.Children == nil {
  3294  		t.Fatalf("nil children summary")
  3295  	}
  3296  	if summary.Children.Pending != 0 || summary.Children.Running != 1 || summary.Children.Dead != 0 {
  3297  		t.Fatalf("bad children summary: %v", summary.Children)
  3298  	}
  3299  
  3300  	// Create watchsets so we can test that update fires the watch
  3301  	ws = memdb.NewWatchSet()
  3302  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  3303  		t.Fatalf("bad: %v", err)
  3304  	}
  3305  
  3306  	// Create the delta updates
  3307  	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStateRunning}}
  3308  	update := &structs.Allocation{
  3309  		ID:           alloc.ID,
  3310  		ClientStatus: structs.AllocClientStatusComplete,
  3311  		TaskStates:   ts,
  3312  		JobID:        alloc.JobID,
  3313  		TaskGroup:    alloc.TaskGroup,
  3314  	}
  3315  	err = state.UpdateAllocsFromClient(1001, []*structs.Allocation{update})
  3316  	if err != nil {
  3317  		t.Fatalf("err: %v", err)
  3318  	}
  3319  
  3320  	if !watchFired(ws) {
  3321  		t.Fatalf("bad")
  3322  	}
  3323  
  3324  	ws = memdb.NewWatchSet()
  3325  	summary, err = state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  3326  	if err != nil {
  3327  		t.Fatalf("err: %v", err)
  3328  	}
  3329  	if summary == nil {
  3330  		t.Fatalf("nil summary")
  3331  	}
  3332  	if summary.JobID != parent.ID {
  3333  		t.Fatalf("bad summary id: %v", parent.ID)
  3334  	}
  3335  	if summary.Children == nil {
  3336  		t.Fatalf("nil children summary")
  3337  	}
  3338  	if summary.Children.Pending != 0 || summary.Children.Running != 0 || summary.Children.Dead != 1 {
  3339  		t.Fatalf("bad children summary: %v", summary.Children)
  3340  	}
  3341  
  3342  	if watchFired(ws) {
  3343  		t.Fatalf("bad")
  3344  	}
  3345  }
  3346  
  3347  func TestStateStore_UpdateAllocsFromClient_ChildJob(t *testing.T) {
  3348  	state := testStateStore(t)
  3349  	alloc1 := mock.Alloc()
  3350  	alloc2 := mock.Alloc()
  3351  
  3352  	if err := state.UpsertJob(999, alloc1.Job); err != nil {
  3353  		t.Fatalf("err: %v", err)
  3354  	}
  3355  	if err := state.UpsertJob(999, alloc2.Job); err != nil {
  3356  		t.Fatalf("err: %v", err)
  3357  	}
  3358  
  3359  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc1, alloc2})
  3360  	if err != nil {
  3361  		t.Fatalf("err: %v", err)
  3362  	}
  3363  
  3364  	// Create watchsets so we can test that update fires the watch
  3365  	watches := make([]memdb.WatchSet, 8)
  3366  	for i := 0; i < 8; i++ {
  3367  		watches[i] = memdb.NewWatchSet()
  3368  	}
  3369  	if _, err := state.AllocByID(watches[0], alloc1.ID); err != nil {
  3370  		t.Fatalf("bad: %v", err)
  3371  	}
  3372  	if _, err := state.AllocByID(watches[1], alloc2.ID); err != nil {
  3373  		t.Fatalf("bad: %v", err)
  3374  	}
  3375  	if _, err := state.AllocsByEval(watches[2], alloc1.EvalID); err != nil {
  3376  		t.Fatalf("bad: %v", err)
  3377  	}
  3378  	if _, err := state.AllocsByEval(watches[3], alloc2.EvalID); err != nil {
  3379  		t.Fatalf("bad: %v", err)
  3380  	}
  3381  	if _, err := state.AllocsByJob(watches[4], alloc1.Namespace, alloc1.JobID, false); err != nil {
  3382  		t.Fatalf("bad: %v", err)
  3383  	}
  3384  	if _, err := state.AllocsByJob(watches[5], alloc2.Namespace, alloc2.JobID, false); err != nil {
  3385  		t.Fatalf("bad: %v", err)
  3386  	}
  3387  	if _, err := state.AllocsByNode(watches[6], alloc1.NodeID); err != nil {
  3388  		t.Fatalf("bad: %v", err)
  3389  	}
  3390  	if _, err := state.AllocsByNode(watches[7], alloc2.NodeID); err != nil {
  3391  		t.Fatalf("bad: %v", err)
  3392  	}
  3393  
  3394  	// Create the delta updates
  3395  	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStatePending}}
  3396  	update := &structs.Allocation{
  3397  		ID:           alloc1.ID,
  3398  		ClientStatus: structs.AllocClientStatusFailed,
  3399  		TaskStates:   ts,
  3400  		JobID:        alloc1.JobID,
  3401  		TaskGroup:    alloc1.TaskGroup,
  3402  	}
  3403  	update2 := &structs.Allocation{
  3404  		ID:           alloc2.ID,
  3405  		ClientStatus: structs.AllocClientStatusRunning,
  3406  		TaskStates:   ts,
  3407  		JobID:        alloc2.JobID,
  3408  		TaskGroup:    alloc2.TaskGroup,
  3409  	}
  3410  
  3411  	err = state.UpdateAllocsFromClient(1001, []*structs.Allocation{update, update2})
  3412  	if err != nil {
  3413  		t.Fatalf("err: %v", err)
  3414  	}
  3415  
  3416  	for i, ws := range watches {
  3417  		if !watchFired(ws) {
  3418  			t.Fatalf("bad %d", i)
  3419  		}
  3420  	}
  3421  
  3422  	ws := memdb.NewWatchSet()
  3423  	out, err := state.AllocByID(ws, alloc1.ID)
  3424  	if err != nil {
  3425  		t.Fatalf("err: %v", err)
  3426  	}
  3427  
  3428  	alloc1.CreateIndex = 1000
  3429  	alloc1.ModifyIndex = 1001
  3430  	alloc1.TaskStates = ts
  3431  	alloc1.ClientStatus = structs.AllocClientStatusFailed
  3432  	if !reflect.DeepEqual(alloc1, out) {
  3433  		t.Fatalf("bad: %#v %#v", alloc1, out)
  3434  	}
  3435  
  3436  	out, err = state.AllocByID(ws, alloc2.ID)
  3437  	if err != nil {
  3438  		t.Fatalf("err: %v", err)
  3439  	}
  3440  
  3441  	alloc2.ModifyIndex = 1000
  3442  	alloc2.ModifyIndex = 1001
  3443  	alloc2.ClientStatus = structs.AllocClientStatusRunning
  3444  	alloc2.TaskStates = ts
  3445  	if !reflect.DeepEqual(alloc2, out) {
  3446  		t.Fatalf("bad: %#v %#v", alloc2, out)
  3447  	}
  3448  
  3449  	index, err := state.Index("allocs")
  3450  	if err != nil {
  3451  		t.Fatalf("err: %v", err)
  3452  	}
  3453  	if index != 1001 {
  3454  		t.Fatalf("bad: %d", index)
  3455  	}
  3456  
  3457  	// Ensure summaries have been updated
  3458  	summary, err := state.JobSummaryByID(ws, alloc1.Namespace, alloc1.JobID)
  3459  	if err != nil {
  3460  		t.Fatalf("err: %v", err)
  3461  	}
  3462  	tgSummary := summary.Summary["web"]
  3463  	if tgSummary.Failed != 1 {
  3464  		t.Fatalf("expected failed: %v, actual: %v, summary: %#v", 1, tgSummary.Failed, tgSummary)
  3465  	}
  3466  
  3467  	summary2, err := state.JobSummaryByID(ws, alloc2.Namespace, alloc2.JobID)
  3468  	if err != nil {
  3469  		t.Fatalf("err: %v", err)
  3470  	}
  3471  	tgSummary2 := summary2.Summary["web"]
  3472  	if tgSummary2.Running != 1 {
  3473  		t.Fatalf("expected running: %v, actual: %v", 1, tgSummary2.Running)
  3474  	}
  3475  
  3476  	if watchFired(ws) {
  3477  		t.Fatalf("bad")
  3478  	}
  3479  }
  3480  
  3481  func TestStateStore_UpdateMultipleAllocsFromClient(t *testing.T) {
  3482  	state := testStateStore(t)
  3483  	alloc := mock.Alloc()
  3484  
  3485  	if err := state.UpsertJob(999, alloc.Job); err != nil {
  3486  		t.Fatalf("err: %v", err)
  3487  	}
  3488  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3489  	if err != nil {
  3490  		t.Fatalf("err: %v", err)
  3491  	}
  3492  
  3493  	// Create the delta updates
  3494  	ts := map[string]*structs.TaskState{"web": {State: structs.TaskStatePending}}
  3495  	update := &structs.Allocation{
  3496  		ID:           alloc.ID,
  3497  		ClientStatus: structs.AllocClientStatusRunning,
  3498  		TaskStates:   ts,
  3499  		JobID:        alloc.JobID,
  3500  		TaskGroup:    alloc.TaskGroup,
  3501  	}
  3502  	update2 := &structs.Allocation{
  3503  		ID:           alloc.ID,
  3504  		ClientStatus: structs.AllocClientStatusPending,
  3505  		TaskStates:   ts,
  3506  		JobID:        alloc.JobID,
  3507  		TaskGroup:    alloc.TaskGroup,
  3508  	}
  3509  
  3510  	err = state.UpdateAllocsFromClient(1001, []*structs.Allocation{update, update2})
  3511  	if err != nil {
  3512  		t.Fatalf("err: %v", err)
  3513  	}
  3514  
  3515  	ws := memdb.NewWatchSet()
  3516  	out, err := state.AllocByID(ws, alloc.ID)
  3517  	if err != nil {
  3518  		t.Fatalf("err: %v", err)
  3519  	}
  3520  
  3521  	alloc.CreateIndex = 1000
  3522  	alloc.ModifyIndex = 1001
  3523  	alloc.TaskStates = ts
  3524  	alloc.ClientStatus = structs.AllocClientStatusPending
  3525  	if !reflect.DeepEqual(alloc, out) {
  3526  		t.Fatalf("bad: %#v , actual:%#v", alloc, out)
  3527  	}
  3528  
  3529  	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  3530  	expectedSummary := &structs.JobSummary{
  3531  		JobID:     alloc.JobID,
  3532  		Namespace: alloc.Namespace,
  3533  		Summary: map[string]structs.TaskGroupSummary{
  3534  			"web": {
  3535  				Starting: 1,
  3536  			},
  3537  		},
  3538  		Children:    new(structs.JobChildrenSummary),
  3539  		CreateIndex: 999,
  3540  		ModifyIndex: 1001,
  3541  	}
  3542  	if err != nil {
  3543  		t.Fatalf("err: %v", err)
  3544  	}
  3545  	if !reflect.DeepEqual(summary, expectedSummary) {
  3546  		t.Fatalf("expected: %#v, actual: %#v", expectedSummary, summary)
  3547  	}
  3548  }
  3549  
  3550  func TestStateStore_UpdateAllocsFromClient_Deployment(t *testing.T) {
  3551  	require := require.New(t)
  3552  	state := testStateStore(t)
  3553  
  3554  	alloc := mock.Alloc()
  3555  	now := time.Now()
  3556  	alloc.CreateTime = now.UnixNano()
  3557  	pdeadline := 5 * time.Minute
  3558  	deployment := mock.Deployment()
  3559  	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
  3560  	alloc.DeploymentID = deployment.ID
  3561  
  3562  	require.Nil(state.UpsertJob(999, alloc.Job))
  3563  	require.Nil(state.UpsertDeployment(1000, deployment))
  3564  	require.Nil(state.UpsertAllocs(1001, []*structs.Allocation{alloc}))
  3565  
  3566  	healthy := now.Add(time.Second)
  3567  	update := &structs.Allocation{
  3568  		ID:           alloc.ID,
  3569  		ClientStatus: structs.AllocClientStatusRunning,
  3570  		JobID:        alloc.JobID,
  3571  		TaskGroup:    alloc.TaskGroup,
  3572  		DeploymentStatus: &structs.AllocDeploymentStatus{
  3573  			Healthy:   helper.BoolToPtr(true),
  3574  			Timestamp: healthy,
  3575  		},
  3576  	}
  3577  	require.Nil(state.UpdateAllocsFromClient(1001, []*structs.Allocation{update}))
  3578  
  3579  	// Check that the deployment state was updated because the healthy
  3580  	// deployment
  3581  	dout, err := state.DeploymentByID(nil, deployment.ID)
  3582  	require.Nil(err)
  3583  	require.NotNil(dout)
  3584  	require.Len(dout.TaskGroups, 1)
  3585  	dstate := dout.TaskGroups[alloc.TaskGroup]
  3586  	require.NotNil(dstate)
  3587  	require.Equal(1, dstate.PlacedAllocs)
  3588  	require.True(healthy.Add(pdeadline).Equal(dstate.RequireProgressBy))
  3589  }
  3590  
  3591  // This tests that the deployment state is merged correctly
  3592  func TestStateStore_UpdateAllocsFromClient_DeploymentStateMerges(t *testing.T) {
  3593  	require := require.New(t)
  3594  	state := testStateStore(t)
  3595  
  3596  	alloc := mock.Alloc()
  3597  	now := time.Now()
  3598  	alloc.CreateTime = now.UnixNano()
  3599  	pdeadline := 5 * time.Minute
  3600  	deployment := mock.Deployment()
  3601  	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
  3602  	alloc.DeploymentID = deployment.ID
  3603  	alloc.DeploymentStatus = &structs.AllocDeploymentStatus{
  3604  		Canary: true,
  3605  	}
  3606  
  3607  	require.Nil(state.UpsertJob(999, alloc.Job))
  3608  	require.Nil(state.UpsertDeployment(1000, deployment))
  3609  	require.Nil(state.UpsertAllocs(1001, []*structs.Allocation{alloc}))
  3610  
  3611  	update := &structs.Allocation{
  3612  		ID:           alloc.ID,
  3613  		ClientStatus: structs.AllocClientStatusRunning,
  3614  		JobID:        alloc.JobID,
  3615  		TaskGroup:    alloc.TaskGroup,
  3616  		DeploymentStatus: &structs.AllocDeploymentStatus{
  3617  			Healthy: helper.BoolToPtr(true),
  3618  			Canary:  false,
  3619  		},
  3620  	}
  3621  	require.Nil(state.UpdateAllocsFromClient(1001, []*structs.Allocation{update}))
  3622  
  3623  	// Check that the merging of the deployment status was correct
  3624  	out, err := state.AllocByID(nil, alloc.ID)
  3625  	require.Nil(err)
  3626  	require.NotNil(out)
  3627  	require.True(out.DeploymentStatus.Canary)
  3628  }
  3629  
  3630  func TestStateStore_UpsertAlloc_Alloc(t *testing.T) {
  3631  	state := testStateStore(t)
  3632  	alloc := mock.Alloc()
  3633  
  3634  	if err := state.UpsertJob(999, alloc.Job); err != nil {
  3635  		t.Fatalf("err: %v", err)
  3636  	}
  3637  
  3638  	// Create watchsets so we can test that update fires the watch
  3639  	watches := make([]memdb.WatchSet, 4)
  3640  	for i := 0; i < 4; i++ {
  3641  		watches[i] = memdb.NewWatchSet()
  3642  	}
  3643  	if _, err := state.AllocByID(watches[0], alloc.ID); err != nil {
  3644  		t.Fatalf("bad: %v", err)
  3645  	}
  3646  	if _, err := state.AllocsByEval(watches[1], alloc.EvalID); err != nil {
  3647  		t.Fatalf("bad: %v", err)
  3648  	}
  3649  	if _, err := state.AllocsByJob(watches[2], alloc.Namespace, alloc.JobID, false); err != nil {
  3650  		t.Fatalf("bad: %v", err)
  3651  	}
  3652  	if _, err := state.AllocsByNode(watches[3], alloc.NodeID); err != nil {
  3653  		t.Fatalf("bad: %v", err)
  3654  	}
  3655  
  3656  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3657  	if err != nil {
  3658  		t.Fatalf("err: %v", err)
  3659  	}
  3660  
  3661  	for i, ws := range watches {
  3662  		if !watchFired(ws) {
  3663  			t.Fatalf("bad %d", i)
  3664  		}
  3665  	}
  3666  
  3667  	ws := memdb.NewWatchSet()
  3668  	out, err := state.AllocByID(ws, alloc.ID)
  3669  	if err != nil {
  3670  		t.Fatalf("err: %v", err)
  3671  	}
  3672  
  3673  	if !reflect.DeepEqual(alloc, out) {
  3674  		t.Fatalf("bad: %#v %#v", alloc, out)
  3675  	}
  3676  
  3677  	index, err := state.Index("allocs")
  3678  	if err != nil {
  3679  		t.Fatalf("err: %v", err)
  3680  	}
  3681  	if index != 1000 {
  3682  		t.Fatalf("bad: %d", index)
  3683  	}
  3684  
  3685  	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  3686  	if err != nil {
  3687  		t.Fatalf("err: %v", err)
  3688  	}
  3689  
  3690  	tgSummary, ok := summary.Summary["web"]
  3691  	if !ok {
  3692  		t.Fatalf("no summary for task group web")
  3693  	}
  3694  	if tgSummary.Starting != 1 {
  3695  		t.Fatalf("expected queued: %v, actual: %v", 1, tgSummary.Starting)
  3696  	}
  3697  
  3698  	if watchFired(ws) {
  3699  		t.Fatalf("bad")
  3700  	}
  3701  }
  3702  
  3703  func TestStateStore_UpsertAlloc_Deployment(t *testing.T) {
  3704  	require := require.New(t)
  3705  	state := testStateStore(t)
  3706  	alloc := mock.Alloc()
  3707  	now := time.Now()
  3708  	alloc.CreateTime = now.UnixNano()
  3709  	alloc.ModifyTime = now.UnixNano()
  3710  	pdeadline := 5 * time.Minute
  3711  	deployment := mock.Deployment()
  3712  	deployment.TaskGroups[alloc.TaskGroup].ProgressDeadline = pdeadline
  3713  	alloc.DeploymentID = deployment.ID
  3714  
  3715  	require.Nil(state.UpsertJob(999, alloc.Job))
  3716  	require.Nil(state.UpsertDeployment(1000, deployment))
  3717  
  3718  	// Create a watch set so we can test that update fires the watch
  3719  	ws := memdb.NewWatchSet()
  3720  	require.Nil(state.AllocsByDeployment(ws, alloc.DeploymentID))
  3721  
  3722  	err := state.UpsertAllocs(1001, []*structs.Allocation{alloc})
  3723  	require.Nil(err)
  3724  
  3725  	if !watchFired(ws) {
  3726  		t.Fatalf("watch not fired")
  3727  	}
  3728  
  3729  	ws = memdb.NewWatchSet()
  3730  	allocs, err := state.AllocsByDeployment(ws, alloc.DeploymentID)
  3731  	require.Nil(err)
  3732  	require.Len(allocs, 1)
  3733  	require.EqualValues(alloc, allocs[0])
  3734  
  3735  	index, err := state.Index("allocs")
  3736  	require.Nil(err)
  3737  	require.EqualValues(1001, index)
  3738  	if watchFired(ws) {
  3739  		t.Fatalf("bad")
  3740  	}
  3741  
  3742  	// Check that the deployment state was updated
  3743  	dout, err := state.DeploymentByID(nil, deployment.ID)
  3744  	require.Nil(err)
  3745  	require.NotNil(dout)
  3746  	require.Len(dout.TaskGroups, 1)
  3747  	dstate := dout.TaskGroups[alloc.TaskGroup]
  3748  	require.NotNil(dstate)
  3749  	require.Equal(1, dstate.PlacedAllocs)
  3750  	require.True(now.Add(pdeadline).Equal(dstate.RequireProgressBy))
  3751  }
  3752  
  3753  // Testing to ensure we keep issue
  3754  // https://github.com/hashicorp/nomad/issues/2583 fixed
  3755  func TestStateStore_UpsertAlloc_No_Job(t *testing.T) {
  3756  	state := testStateStore(t)
  3757  	alloc := mock.Alloc()
  3758  	alloc.Job = nil
  3759  
  3760  	err := state.UpsertAllocs(999, []*structs.Allocation{alloc})
  3761  	if err == nil || !strings.Contains(err.Error(), "without a job") {
  3762  		t.Fatalf("expect err: %v", err)
  3763  	}
  3764  }
  3765  
  3766  func TestStateStore_UpsertAlloc_NoEphemeralDisk(t *testing.T) {
  3767  	state := testStateStore(t)
  3768  	alloc := mock.Alloc()
  3769  	alloc.Job.TaskGroups[0].EphemeralDisk = nil
  3770  	alloc.Job.TaskGroups[0].Tasks[0].Resources.DiskMB = 120
  3771  
  3772  	if err := state.UpsertJob(999, alloc.Job); err != nil {
  3773  		t.Fatalf("err: %v", err)
  3774  	}
  3775  
  3776  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3777  	if err != nil {
  3778  		t.Fatalf("err: %v", err)
  3779  	}
  3780  
  3781  	ws := memdb.NewWatchSet()
  3782  	out, err := state.AllocByID(ws, alloc.ID)
  3783  	if err != nil {
  3784  		t.Fatalf("err: %v", err)
  3785  	}
  3786  
  3787  	expected := alloc.Copy()
  3788  	expected.Job.TaskGroups[0].EphemeralDisk = &structs.EphemeralDisk{SizeMB: 120}
  3789  	if !reflect.DeepEqual(expected, out) {
  3790  		t.Fatalf("bad: %#v %#v", expected, out)
  3791  	}
  3792  }
  3793  
  3794  func TestStateStore_UpsertAlloc_ChildJob(t *testing.T) {
  3795  	state := testStateStore(t)
  3796  
  3797  	parent := mock.Job()
  3798  	if err := state.UpsertJob(998, parent); err != nil {
  3799  		t.Fatalf("err: %v", err)
  3800  	}
  3801  
  3802  	child := mock.Job()
  3803  	child.ParentID = parent.ID
  3804  
  3805  	if err := state.UpsertJob(999, child); err != nil {
  3806  		t.Fatalf("err: %v", err)
  3807  	}
  3808  
  3809  	alloc := mock.Alloc()
  3810  	alloc.JobID = child.ID
  3811  	alloc.Job = child
  3812  
  3813  	// Create watchsets so we can test that delete fires the watch
  3814  	ws := memdb.NewWatchSet()
  3815  	if _, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID); err != nil {
  3816  		t.Fatalf("bad: %v", err)
  3817  	}
  3818  
  3819  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3820  	if err != nil {
  3821  		t.Fatalf("err: %v", err)
  3822  	}
  3823  
  3824  	if !watchFired(ws) {
  3825  		t.Fatalf("bad")
  3826  	}
  3827  
  3828  	ws = memdb.NewWatchSet()
  3829  	summary, err := state.JobSummaryByID(ws, parent.Namespace, parent.ID)
  3830  	if err != nil {
  3831  		t.Fatalf("err: %v", err)
  3832  	}
  3833  	if summary == nil {
  3834  		t.Fatalf("nil summary")
  3835  	}
  3836  	if summary.JobID != parent.ID {
  3837  		t.Fatalf("bad summary id: %v", parent.ID)
  3838  	}
  3839  	if summary.Children == nil {
  3840  		t.Fatalf("nil children summary")
  3841  	}
  3842  	if summary.Children.Pending != 0 || summary.Children.Running != 1 || summary.Children.Dead != 0 {
  3843  		t.Fatalf("bad children summary: %v", summary.Children)
  3844  	}
  3845  
  3846  	if watchFired(ws) {
  3847  		t.Fatalf("bad")
  3848  	}
  3849  }
  3850  
  3851  func TestStateStore_UpdateAlloc_Alloc(t *testing.T) {
  3852  	state := testStateStore(t)
  3853  	alloc := mock.Alloc()
  3854  
  3855  	if err := state.UpsertJob(999, alloc.Job); err != nil {
  3856  		t.Fatalf("err: %v", err)
  3857  	}
  3858  
  3859  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3860  	if err != nil {
  3861  		t.Fatalf("err: %v", err)
  3862  	}
  3863  
  3864  	ws := memdb.NewWatchSet()
  3865  	summary, err := state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  3866  	if err != nil {
  3867  		t.Fatalf("err: %v", err)
  3868  	}
  3869  	tgSummary := summary.Summary["web"]
  3870  	if tgSummary.Starting != 1 {
  3871  		t.Fatalf("expected starting: %v, actual: %v", 1, tgSummary.Starting)
  3872  	}
  3873  
  3874  	alloc2 := mock.Alloc()
  3875  	alloc2.ID = alloc.ID
  3876  	alloc2.NodeID = alloc.NodeID + ".new"
  3877  	state.UpsertJobSummary(1001, mock.JobSummary(alloc2.JobID))
  3878  
  3879  	// Create watchsets so we can test that update fires the watch
  3880  	watches := make([]memdb.WatchSet, 4)
  3881  	for i := 0; i < 4; i++ {
  3882  		watches[i] = memdb.NewWatchSet()
  3883  	}
  3884  	if _, err := state.AllocByID(watches[0], alloc2.ID); err != nil {
  3885  		t.Fatalf("bad: %v", err)
  3886  	}
  3887  	if _, err := state.AllocsByEval(watches[1], alloc2.EvalID); err != nil {
  3888  		t.Fatalf("bad: %v", err)
  3889  	}
  3890  	if _, err := state.AllocsByJob(watches[2], alloc2.Namespace, alloc2.JobID, false); err != nil {
  3891  		t.Fatalf("bad: %v", err)
  3892  	}
  3893  	if _, err := state.AllocsByNode(watches[3], alloc2.NodeID); err != nil {
  3894  		t.Fatalf("bad: %v", err)
  3895  	}
  3896  
  3897  	err = state.UpsertAllocs(1002, []*structs.Allocation{alloc2})
  3898  	if err != nil {
  3899  		t.Fatalf("err: %v", err)
  3900  	}
  3901  
  3902  	for i, ws := range watches {
  3903  		if !watchFired(ws) {
  3904  			t.Fatalf("bad %d", i)
  3905  		}
  3906  	}
  3907  
  3908  	ws = memdb.NewWatchSet()
  3909  	out, err := state.AllocByID(ws, alloc.ID)
  3910  	if err != nil {
  3911  		t.Fatalf("err: %v", err)
  3912  	}
  3913  
  3914  	if !reflect.DeepEqual(alloc2, out) {
  3915  		t.Fatalf("bad: %#v %#v", alloc2, out)
  3916  	}
  3917  
  3918  	if out.CreateIndex != 1000 {
  3919  		t.Fatalf("bad: %#v", out)
  3920  	}
  3921  	if out.ModifyIndex != 1002 {
  3922  		t.Fatalf("bad: %#v", out)
  3923  	}
  3924  
  3925  	index, err := state.Index("allocs")
  3926  	if err != nil {
  3927  		t.Fatalf("err: %v", err)
  3928  	}
  3929  	if index != 1002 {
  3930  		t.Fatalf("bad: %d", index)
  3931  	}
  3932  
  3933  	// Ensure that summary hasb't changed
  3934  	summary, err = state.JobSummaryByID(ws, alloc.Namespace, alloc.JobID)
  3935  	if err != nil {
  3936  		t.Fatalf("err: %v", err)
  3937  	}
  3938  	tgSummary = summary.Summary["web"]
  3939  	if tgSummary.Starting != 1 {
  3940  		t.Fatalf("expected starting: %v, actual: %v", 1, tgSummary.Starting)
  3941  	}
  3942  
  3943  	if watchFired(ws) {
  3944  		t.Fatalf("bad")
  3945  	}
  3946  }
  3947  
  3948  // This test ensures that the state store will mark the clients status as lost
  3949  // when set rather than preferring the existing status.
  3950  func TestStateStore_UpdateAlloc_Lost(t *testing.T) {
  3951  	state := testStateStore(t)
  3952  	alloc := mock.Alloc()
  3953  	alloc.ClientStatus = "foo"
  3954  
  3955  	if err := state.UpsertJob(999, alloc.Job); err != nil {
  3956  		t.Fatalf("err: %v", err)
  3957  	}
  3958  
  3959  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3960  	if err != nil {
  3961  		t.Fatalf("err: %v", err)
  3962  	}
  3963  
  3964  	alloc2 := new(structs.Allocation)
  3965  	*alloc2 = *alloc
  3966  	alloc2.ClientStatus = structs.AllocClientStatusLost
  3967  	if err := state.UpsertAllocs(1001, []*structs.Allocation{alloc2}); err != nil {
  3968  		t.Fatalf("err: %v", err)
  3969  	}
  3970  
  3971  	ws := memdb.NewWatchSet()
  3972  	out, err := state.AllocByID(ws, alloc2.ID)
  3973  	if err != nil {
  3974  		t.Fatalf("err: %v", err)
  3975  	}
  3976  
  3977  	if out.ClientStatus != structs.AllocClientStatusLost {
  3978  		t.Fatalf("bad: %#v", out)
  3979  	}
  3980  }
  3981  
  3982  // This test ensures an allocation can be updated when there is no job
  3983  // associated with it. This will happen when a job is stopped by an user which
  3984  // has non-terminal allocations on clients
  3985  func TestStateStore_UpdateAlloc_NoJob(t *testing.T) {
  3986  	state := testStateStore(t)
  3987  	alloc := mock.Alloc()
  3988  
  3989  	// Upsert a job
  3990  	state.UpsertJobSummary(998, mock.JobSummary(alloc.JobID))
  3991  	if err := state.UpsertJob(999, alloc.Job); err != nil {
  3992  		t.Fatalf("err: %v", err)
  3993  	}
  3994  
  3995  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  3996  	if err != nil {
  3997  		t.Fatalf("err: %v", err)
  3998  	}
  3999  
  4000  	if err := state.DeleteJob(1001, alloc.Namespace, alloc.JobID); err != nil {
  4001  		t.Fatalf("err: %v", err)
  4002  	}
  4003  
  4004  	// Update the desired state of the allocation to stop
  4005  	allocCopy := alloc.Copy()
  4006  	allocCopy.DesiredStatus = structs.AllocDesiredStatusStop
  4007  	if err := state.UpsertAllocs(1002, []*structs.Allocation{allocCopy}); err != nil {
  4008  		t.Fatalf("err: %v", err)
  4009  	}
  4010  
  4011  	// Update the client state of the allocation to complete
  4012  	allocCopy1 := allocCopy.Copy()
  4013  	allocCopy1.ClientStatus = structs.AllocClientStatusComplete
  4014  	if err := state.UpdateAllocsFromClient(1003, []*structs.Allocation{allocCopy1}); err != nil {
  4015  		t.Fatalf("err: %v", err)
  4016  	}
  4017  
  4018  	ws := memdb.NewWatchSet()
  4019  	out, _ := state.AllocByID(ws, alloc.ID)
  4020  	// Update the modify index of the alloc before comparing
  4021  	allocCopy1.ModifyIndex = 1003
  4022  	if !reflect.DeepEqual(out, allocCopy1) {
  4023  		t.Fatalf("expected: %#v \n actual: %#v", allocCopy1, out)
  4024  	}
  4025  }
  4026  
  4027  func TestStateStore_UpdateAllocDesiredTransition(t *testing.T) {
  4028  	t.Parallel()
  4029  	require := require.New(t)
  4030  
  4031  	state := testStateStore(t)
  4032  	alloc := mock.Alloc()
  4033  
  4034  	require.Nil(state.UpsertJob(999, alloc.Job))
  4035  	require.Nil(state.UpsertAllocs(1000, []*structs.Allocation{alloc}))
  4036  
  4037  	t1 := &structs.DesiredTransition{
  4038  		Migrate: helper.BoolToPtr(true),
  4039  	}
  4040  	t2 := &structs.DesiredTransition{
  4041  		Migrate: helper.BoolToPtr(false),
  4042  	}
  4043  	eval := &structs.Evaluation{
  4044  		ID:             uuid.Generate(),
  4045  		Namespace:      alloc.Namespace,
  4046  		Priority:       alloc.Job.Priority,
  4047  		Type:           alloc.Job.Type,
  4048  		TriggeredBy:    structs.EvalTriggerNodeDrain,
  4049  		JobID:          alloc.Job.ID,
  4050  		JobModifyIndex: alloc.Job.ModifyIndex,
  4051  		Status:         structs.EvalStatusPending,
  4052  	}
  4053  	evals := []*structs.Evaluation{eval}
  4054  
  4055  	m := map[string]*structs.DesiredTransition{alloc.ID: t1}
  4056  	require.Nil(state.UpdateAllocsDesiredTransitions(1001, m, evals))
  4057  
  4058  	ws := memdb.NewWatchSet()
  4059  	out, err := state.AllocByID(ws, alloc.ID)
  4060  	require.Nil(err)
  4061  	require.NotNil(out.DesiredTransition.Migrate)
  4062  	require.True(*out.DesiredTransition.Migrate)
  4063  	require.EqualValues(1000, out.CreateIndex)
  4064  	require.EqualValues(1001, out.ModifyIndex)
  4065  
  4066  	index, err := state.Index("allocs")
  4067  	require.Nil(err)
  4068  	require.EqualValues(1001, index)
  4069  
  4070  	// Check the eval is created
  4071  	eout, err := state.EvalByID(nil, eval.ID)
  4072  	require.Nil(err)
  4073  	require.NotNil(eout)
  4074  
  4075  	m = map[string]*structs.DesiredTransition{alloc.ID: t2}
  4076  	require.Nil(state.UpdateAllocsDesiredTransitions(1002, m, evals))
  4077  
  4078  	ws = memdb.NewWatchSet()
  4079  	out, err = state.AllocByID(ws, alloc.ID)
  4080  	require.Nil(err)
  4081  	require.NotNil(out.DesiredTransition.Migrate)
  4082  	require.False(*out.DesiredTransition.Migrate)
  4083  	require.EqualValues(1000, out.CreateIndex)
  4084  	require.EqualValues(1002, out.ModifyIndex)
  4085  
  4086  	index, err = state.Index("allocs")
  4087  	require.Nil(err)
  4088  	require.EqualValues(1002, index)
  4089  
  4090  	// Try with a bogus alloc id
  4091  	m = map[string]*structs.DesiredTransition{uuid.Generate(): t2}
  4092  	require.Nil(state.UpdateAllocsDesiredTransitions(1003, m, evals))
  4093  }
  4094  
  4095  func TestStateStore_JobSummary(t *testing.T) {
  4096  	state := testStateStore(t)
  4097  
  4098  	// Add a job
  4099  	job := mock.Job()
  4100  	state.UpsertJob(900, job)
  4101  
  4102  	// Get the job back
  4103  	ws := memdb.NewWatchSet()
  4104  	outJob, _ := state.JobByID(ws, job.Namespace, job.ID)
  4105  	if outJob.CreateIndex != 900 {
  4106  		t.Fatalf("bad create index: %v", outJob.CreateIndex)
  4107  	}
  4108  	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
  4109  	if summary.CreateIndex != 900 {
  4110  		t.Fatalf("bad create index: %v", summary.CreateIndex)
  4111  	}
  4112  
  4113  	// Upsert an allocation
  4114  	alloc := mock.Alloc()
  4115  	alloc.JobID = job.ID
  4116  	alloc.Job = job
  4117  	state.UpsertAllocs(910, []*structs.Allocation{alloc})
  4118  
  4119  	// Update the alloc from client
  4120  	alloc1 := alloc.Copy()
  4121  	alloc1.ClientStatus = structs.AllocClientStatusPending
  4122  	alloc1.DesiredStatus = ""
  4123  	state.UpdateAllocsFromClient(920, []*structs.Allocation{alloc})
  4124  
  4125  	alloc3 := alloc.Copy()
  4126  	alloc3.ClientStatus = structs.AllocClientStatusRunning
  4127  	alloc3.DesiredStatus = ""
  4128  	state.UpdateAllocsFromClient(930, []*structs.Allocation{alloc3})
  4129  
  4130  	// Upsert the alloc
  4131  	alloc4 := alloc.Copy()
  4132  	alloc4.ClientStatus = structs.AllocClientStatusPending
  4133  	alloc4.DesiredStatus = structs.AllocDesiredStatusRun
  4134  	state.UpsertAllocs(950, []*structs.Allocation{alloc4})
  4135  
  4136  	// Again upsert the alloc
  4137  	alloc5 := alloc.Copy()
  4138  	alloc5.ClientStatus = structs.AllocClientStatusPending
  4139  	alloc5.DesiredStatus = structs.AllocDesiredStatusRun
  4140  	state.UpsertAllocs(970, []*structs.Allocation{alloc5})
  4141  
  4142  	if !watchFired(ws) {
  4143  		t.Fatalf("bad")
  4144  	}
  4145  
  4146  	expectedSummary := structs.JobSummary{
  4147  		JobID:     job.ID,
  4148  		Namespace: job.Namespace,
  4149  		Summary: map[string]structs.TaskGroupSummary{
  4150  			"web": {
  4151  				Running: 1,
  4152  			},
  4153  		},
  4154  		Children:    new(structs.JobChildrenSummary),
  4155  		CreateIndex: 900,
  4156  		ModifyIndex: 930,
  4157  	}
  4158  
  4159  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  4160  	if !reflect.DeepEqual(&expectedSummary, summary) {
  4161  		t.Fatalf("expected: %#v, actual: %v", expectedSummary, summary)
  4162  	}
  4163  
  4164  	// De-register the job.
  4165  	state.DeleteJob(980, job.Namespace, job.ID)
  4166  
  4167  	// Shouldn't have any effect on the summary
  4168  	alloc6 := alloc.Copy()
  4169  	alloc6.ClientStatus = structs.AllocClientStatusRunning
  4170  	alloc6.DesiredStatus = ""
  4171  	state.UpdateAllocsFromClient(990, []*structs.Allocation{alloc6})
  4172  
  4173  	// We shouldn't have any summary at this point
  4174  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  4175  	if summary != nil {
  4176  		t.Fatalf("expected nil, actual: %#v", summary)
  4177  	}
  4178  
  4179  	// Re-register the same job
  4180  	job1 := mock.Job()
  4181  	job1.ID = job.ID
  4182  	state.UpsertJob(1000, job1)
  4183  	outJob2, _ := state.JobByID(ws, job1.Namespace, job1.ID)
  4184  	if outJob2.CreateIndex != 1000 {
  4185  		t.Fatalf("bad create index: %v", outJob2.CreateIndex)
  4186  	}
  4187  	summary, _ = state.JobSummaryByID(ws, job1.Namespace, job1.ID)
  4188  	if summary.CreateIndex != 1000 {
  4189  		t.Fatalf("bad create index: %v", summary.CreateIndex)
  4190  	}
  4191  
  4192  	// Upsert an allocation
  4193  	alloc7 := alloc.Copy()
  4194  	alloc7.JobID = outJob.ID
  4195  	alloc7.Job = outJob
  4196  	alloc7.ClientStatus = structs.AllocClientStatusComplete
  4197  	alloc7.DesiredStatus = structs.AllocDesiredStatusRun
  4198  	state.UpdateAllocsFromClient(1020, []*structs.Allocation{alloc7})
  4199  
  4200  	expectedSummary = structs.JobSummary{
  4201  		JobID:     job.ID,
  4202  		Namespace: job.Namespace,
  4203  		Summary: map[string]structs.TaskGroupSummary{
  4204  			"web": {},
  4205  		},
  4206  		Children:    new(structs.JobChildrenSummary),
  4207  		CreateIndex: 1000,
  4208  		ModifyIndex: 1000,
  4209  	}
  4210  
  4211  	summary, _ = state.JobSummaryByID(ws, job1.Namespace, job1.ID)
  4212  	if !reflect.DeepEqual(&expectedSummary, summary) {
  4213  		t.Fatalf("expected: %#v, actual: %#v", expectedSummary, summary)
  4214  	}
  4215  }
  4216  
  4217  func TestStateStore_ReconcileJobSummary(t *testing.T) {
  4218  	state := testStateStore(t)
  4219  
  4220  	// Create an alloc
  4221  	alloc := mock.Alloc()
  4222  
  4223  	// Add another task group to the job
  4224  	tg2 := alloc.Job.TaskGroups[0].Copy()
  4225  	tg2.Name = "db"
  4226  	alloc.Job.TaskGroups = append(alloc.Job.TaskGroups, tg2)
  4227  	state.UpsertJob(100, alloc.Job)
  4228  
  4229  	// Create one more alloc for the db task group
  4230  	alloc2 := mock.Alloc()
  4231  	alloc2.TaskGroup = "db"
  4232  	alloc2.JobID = alloc.JobID
  4233  	alloc2.Job = alloc.Job
  4234  
  4235  	// Upserts the alloc
  4236  	state.UpsertAllocs(110, []*structs.Allocation{alloc, alloc2})
  4237  
  4238  	// Change the state of the first alloc to running
  4239  	alloc3 := alloc.Copy()
  4240  	alloc3.ClientStatus = structs.AllocClientStatusRunning
  4241  	state.UpdateAllocsFromClient(120, []*structs.Allocation{alloc3})
  4242  
  4243  	//Add some more allocs to the second tg
  4244  	alloc4 := mock.Alloc()
  4245  	alloc4.JobID = alloc.JobID
  4246  	alloc4.Job = alloc.Job
  4247  	alloc4.TaskGroup = "db"
  4248  	alloc5 := alloc4.Copy()
  4249  	alloc5.ClientStatus = structs.AllocClientStatusRunning
  4250  
  4251  	alloc6 := mock.Alloc()
  4252  	alloc6.JobID = alloc.JobID
  4253  	alloc6.Job = alloc.Job
  4254  	alloc6.TaskGroup = "db"
  4255  	alloc7 := alloc6.Copy()
  4256  	alloc7.ClientStatus = structs.AllocClientStatusComplete
  4257  
  4258  	alloc8 := mock.Alloc()
  4259  	alloc8.JobID = alloc.JobID
  4260  	alloc8.Job = alloc.Job
  4261  	alloc8.TaskGroup = "db"
  4262  	alloc9 := alloc8.Copy()
  4263  	alloc9.ClientStatus = structs.AllocClientStatusFailed
  4264  
  4265  	alloc10 := mock.Alloc()
  4266  	alloc10.JobID = alloc.JobID
  4267  	alloc10.Job = alloc.Job
  4268  	alloc10.TaskGroup = "db"
  4269  	alloc11 := alloc10.Copy()
  4270  	alloc11.ClientStatus = structs.AllocClientStatusLost
  4271  
  4272  	state.UpsertAllocs(130, []*structs.Allocation{alloc4, alloc6, alloc8, alloc10})
  4273  
  4274  	state.UpdateAllocsFromClient(150, []*structs.Allocation{alloc5, alloc7, alloc9, alloc11})
  4275  
  4276  	// DeleteJobSummary is a helper method and doesn't modify the indexes table
  4277  	state.DeleteJobSummary(130, alloc.Namespace, alloc.Job.ID)
  4278  
  4279  	state.ReconcileJobSummaries(120)
  4280  
  4281  	ws := memdb.NewWatchSet()
  4282  	summary, _ := state.JobSummaryByID(ws, alloc.Namespace, alloc.Job.ID)
  4283  	expectedSummary := structs.JobSummary{
  4284  		JobID:     alloc.Job.ID,
  4285  		Namespace: alloc.Namespace,
  4286  		Summary: map[string]structs.TaskGroupSummary{
  4287  			"web": {
  4288  				Running: 1,
  4289  			},
  4290  			"db": {
  4291  				Starting: 1,
  4292  				Running:  1,
  4293  				Failed:   1,
  4294  				Complete: 1,
  4295  				Lost:     1,
  4296  			},
  4297  		},
  4298  		CreateIndex: 100,
  4299  		ModifyIndex: 120,
  4300  	}
  4301  	if !reflect.DeepEqual(&expectedSummary, summary) {
  4302  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  4303  	}
  4304  }
  4305  
  4306  func TestStateStore_UpdateAlloc_JobNotPresent(t *testing.T) {
  4307  	state := testStateStore(t)
  4308  
  4309  	alloc := mock.Alloc()
  4310  	state.UpsertJob(100, alloc.Job)
  4311  	state.UpsertAllocs(200, []*structs.Allocation{alloc})
  4312  
  4313  	// Delete the job
  4314  	state.DeleteJob(300, alloc.Namespace, alloc.Job.ID)
  4315  
  4316  	// Update the alloc
  4317  	alloc1 := alloc.Copy()
  4318  	alloc1.ClientStatus = structs.AllocClientStatusRunning
  4319  
  4320  	// Updating allocation should not throw any error
  4321  	if err := state.UpdateAllocsFromClient(400, []*structs.Allocation{alloc1}); err != nil {
  4322  		t.Fatalf("expect err: %v", err)
  4323  	}
  4324  
  4325  	// Re-Register the job
  4326  	state.UpsertJob(500, alloc.Job)
  4327  
  4328  	// Update the alloc again
  4329  	alloc2 := alloc.Copy()
  4330  	alloc2.ClientStatus = structs.AllocClientStatusComplete
  4331  	if err := state.UpdateAllocsFromClient(400, []*structs.Allocation{alloc1}); err != nil {
  4332  		t.Fatalf("expect err: %v", err)
  4333  	}
  4334  
  4335  	// Job Summary of the newly registered job shouldn't account for the
  4336  	// allocation update for the older job
  4337  	expectedSummary := structs.JobSummary{
  4338  		JobID:     alloc1.JobID,
  4339  		Namespace: alloc1.Namespace,
  4340  		Summary: map[string]structs.TaskGroupSummary{
  4341  			"web": {},
  4342  		},
  4343  		Children:    new(structs.JobChildrenSummary),
  4344  		CreateIndex: 500,
  4345  		ModifyIndex: 500,
  4346  	}
  4347  
  4348  	ws := memdb.NewWatchSet()
  4349  	summary, _ := state.JobSummaryByID(ws, alloc.Namespace, alloc.Job.ID)
  4350  	if !reflect.DeepEqual(&expectedSummary, summary) {
  4351  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  4352  	}
  4353  }
  4354  
  4355  func TestStateStore_EvictAlloc_Alloc(t *testing.T) {
  4356  	state := testStateStore(t)
  4357  	alloc := mock.Alloc()
  4358  
  4359  	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
  4360  	err := state.UpsertAllocs(1000, []*structs.Allocation{alloc})
  4361  	if err != nil {
  4362  		t.Fatalf("err: %v", err)
  4363  	}
  4364  
  4365  	alloc2 := new(structs.Allocation)
  4366  	*alloc2 = *alloc
  4367  	alloc2.DesiredStatus = structs.AllocDesiredStatusEvict
  4368  	err = state.UpsertAllocs(1001, []*structs.Allocation{alloc2})
  4369  	if err != nil {
  4370  		t.Fatalf("err: %v", err)
  4371  	}
  4372  
  4373  	ws := memdb.NewWatchSet()
  4374  	out, err := state.AllocByID(ws, alloc.ID)
  4375  	if err != nil {
  4376  		t.Fatalf("err: %v", err)
  4377  	}
  4378  
  4379  	if out.DesiredStatus != structs.AllocDesiredStatusEvict {
  4380  		t.Fatalf("bad: %#v %#v", alloc, out)
  4381  	}
  4382  
  4383  	index, err := state.Index("allocs")
  4384  	if err != nil {
  4385  		t.Fatalf("err: %v", err)
  4386  	}
  4387  	if index != 1001 {
  4388  		t.Fatalf("bad: %d", index)
  4389  	}
  4390  }
  4391  
  4392  func TestStateStore_AllocsByNode(t *testing.T) {
  4393  	state := testStateStore(t)
  4394  	var allocs []*structs.Allocation
  4395  
  4396  	for i := 0; i < 10; i++ {
  4397  		alloc := mock.Alloc()
  4398  		alloc.NodeID = "foo"
  4399  		allocs = append(allocs, alloc)
  4400  	}
  4401  
  4402  	for idx, alloc := range allocs {
  4403  		state.UpsertJobSummary(uint64(900+idx), mock.JobSummary(alloc.JobID))
  4404  	}
  4405  
  4406  	err := state.UpsertAllocs(1000, allocs)
  4407  	if err != nil {
  4408  		t.Fatalf("err: %v", err)
  4409  	}
  4410  
  4411  	ws := memdb.NewWatchSet()
  4412  	out, err := state.AllocsByNode(ws, "foo")
  4413  	if err != nil {
  4414  		t.Fatalf("err: %v", err)
  4415  	}
  4416  
  4417  	sort.Sort(AllocIDSort(allocs))
  4418  	sort.Sort(AllocIDSort(out))
  4419  
  4420  	if !reflect.DeepEqual(allocs, out) {
  4421  		t.Fatalf("bad: %#v %#v", allocs, out)
  4422  	}
  4423  
  4424  	if watchFired(ws) {
  4425  		t.Fatalf("bad")
  4426  	}
  4427  }
  4428  
  4429  func TestStateStore_AllocsByNodeTerminal(t *testing.T) {
  4430  	state := testStateStore(t)
  4431  	var allocs, term, nonterm []*structs.Allocation
  4432  
  4433  	for i := 0; i < 10; i++ {
  4434  		alloc := mock.Alloc()
  4435  		alloc.NodeID = "foo"
  4436  		if i%2 == 0 {
  4437  			alloc.DesiredStatus = structs.AllocDesiredStatusStop
  4438  			term = append(term, alloc)
  4439  		} else {
  4440  			nonterm = append(nonterm, alloc)
  4441  		}
  4442  		allocs = append(allocs, alloc)
  4443  	}
  4444  
  4445  	for idx, alloc := range allocs {
  4446  		state.UpsertJobSummary(uint64(900+idx), mock.JobSummary(alloc.JobID))
  4447  	}
  4448  
  4449  	err := state.UpsertAllocs(1000, allocs)
  4450  	if err != nil {
  4451  		t.Fatalf("err: %v", err)
  4452  	}
  4453  
  4454  	// Verify the terminal allocs
  4455  	ws := memdb.NewWatchSet()
  4456  	out, err := state.AllocsByNodeTerminal(ws, "foo", true)
  4457  	if err != nil {
  4458  		t.Fatalf("err: %v", err)
  4459  	}
  4460  
  4461  	sort.Sort(AllocIDSort(term))
  4462  	sort.Sort(AllocIDSort(out))
  4463  
  4464  	if !reflect.DeepEqual(term, out) {
  4465  		t.Fatalf("bad: %#v %#v", term, out)
  4466  	}
  4467  
  4468  	// Verify the non-terminal allocs
  4469  	out, err = state.AllocsByNodeTerminal(ws, "foo", false)
  4470  	if err != nil {
  4471  		t.Fatalf("err: %v", err)
  4472  	}
  4473  
  4474  	sort.Sort(AllocIDSort(nonterm))
  4475  	sort.Sort(AllocIDSort(out))
  4476  
  4477  	if !reflect.DeepEqual(nonterm, out) {
  4478  		t.Fatalf("bad: %#v %#v", nonterm, out)
  4479  	}
  4480  
  4481  	if watchFired(ws) {
  4482  		t.Fatalf("bad")
  4483  	}
  4484  }
  4485  
  4486  func TestStateStore_AllocsByJob(t *testing.T) {
  4487  	state := testStateStore(t)
  4488  	var allocs []*structs.Allocation
  4489  
  4490  	for i := 0; i < 10; i++ {
  4491  		alloc := mock.Alloc()
  4492  		alloc.JobID = "foo"
  4493  		allocs = append(allocs, alloc)
  4494  	}
  4495  
  4496  	for i, alloc := range allocs {
  4497  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  4498  	}
  4499  
  4500  	err := state.UpsertAllocs(1000, allocs)
  4501  	if err != nil {
  4502  		t.Fatalf("err: %v", err)
  4503  	}
  4504  
  4505  	ws := memdb.NewWatchSet()
  4506  	out, err := state.AllocsByJob(ws, mock.Alloc().Namespace, "foo", false)
  4507  	if err != nil {
  4508  		t.Fatalf("err: %v", err)
  4509  	}
  4510  
  4511  	sort.Sort(AllocIDSort(allocs))
  4512  	sort.Sort(AllocIDSort(out))
  4513  
  4514  	if !reflect.DeepEqual(allocs, out) {
  4515  		t.Fatalf("bad: %#v %#v", allocs, out)
  4516  	}
  4517  
  4518  	if watchFired(ws) {
  4519  		t.Fatalf("bad")
  4520  	}
  4521  }
  4522  
  4523  func TestStateStore_AllocsForRegisteredJob(t *testing.T) {
  4524  	state := testStateStore(t)
  4525  	var allocs []*structs.Allocation
  4526  	var allocs1 []*structs.Allocation
  4527  
  4528  	job := mock.Job()
  4529  	job.ID = "foo"
  4530  	state.UpsertJob(100, job)
  4531  	for i := 0; i < 3; i++ {
  4532  		alloc := mock.Alloc()
  4533  		alloc.Job = job
  4534  		alloc.JobID = job.ID
  4535  		allocs = append(allocs, alloc)
  4536  	}
  4537  	if err := state.UpsertAllocs(200, allocs); err != nil {
  4538  		t.Fatalf("err: %v", err)
  4539  	}
  4540  
  4541  	if err := state.DeleteJob(250, job.Namespace, job.ID); err != nil {
  4542  		t.Fatalf("err: %v", err)
  4543  	}
  4544  
  4545  	job1 := mock.Job()
  4546  	job1.ID = "foo"
  4547  	job1.CreateIndex = 50
  4548  	state.UpsertJob(300, job1)
  4549  	for i := 0; i < 4; i++ {
  4550  		alloc := mock.Alloc()
  4551  		alloc.Job = job1
  4552  		alloc.JobID = job1.ID
  4553  		allocs1 = append(allocs1, alloc)
  4554  	}
  4555  
  4556  	if err := state.UpsertAllocs(1000, allocs1); err != nil {
  4557  		t.Fatalf("err: %v", err)
  4558  	}
  4559  
  4560  	ws := memdb.NewWatchSet()
  4561  	out, err := state.AllocsByJob(ws, job1.Namespace, job1.ID, true)
  4562  	if err != nil {
  4563  		t.Fatalf("err: %v", err)
  4564  	}
  4565  
  4566  	expected := len(allocs) + len(allocs1)
  4567  	if len(out) != expected {
  4568  		t.Fatalf("expected: %v, actual: %v", expected, len(out))
  4569  	}
  4570  
  4571  	out1, err := state.AllocsByJob(ws, job1.Namespace, job1.ID, false)
  4572  	if err != nil {
  4573  		t.Fatalf("bad: %v", err)
  4574  	}
  4575  
  4576  	expected = len(allocs1)
  4577  	if len(out1) != expected {
  4578  		t.Fatalf("expected: %v, actual: %v", expected, len(out1))
  4579  	}
  4580  
  4581  	if watchFired(ws) {
  4582  		t.Fatalf("bad")
  4583  	}
  4584  }
  4585  
  4586  func TestStateStore_AllocsByIDPrefix(t *testing.T) {
  4587  	state := testStateStore(t)
  4588  	var allocs []*structs.Allocation
  4589  
  4590  	ids := []string{
  4591  		"aaaaaaaa-7bfb-395d-eb95-0685af2176b2",
  4592  		"aaaaaaab-7bfb-395d-eb95-0685af2176b2",
  4593  		"aaaaaabb-7bfb-395d-eb95-0685af2176b2",
  4594  		"aaaaabbb-7bfb-395d-eb95-0685af2176b2",
  4595  		"aaaabbbb-7bfb-395d-eb95-0685af2176b2",
  4596  		"aaabbbbb-7bfb-395d-eb95-0685af2176b2",
  4597  		"aabbbbbb-7bfb-395d-eb95-0685af2176b2",
  4598  		"abbbbbbb-7bfb-395d-eb95-0685af2176b2",
  4599  		"bbbbbbbb-7bfb-395d-eb95-0685af2176b2",
  4600  	}
  4601  	for i := 0; i < 9; i++ {
  4602  		alloc := mock.Alloc()
  4603  		alloc.ID = ids[i]
  4604  		allocs = append(allocs, alloc)
  4605  	}
  4606  
  4607  	for i, alloc := range allocs {
  4608  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  4609  	}
  4610  
  4611  	err := state.UpsertAllocs(1000, allocs)
  4612  	if err != nil {
  4613  		t.Fatalf("err: %v", err)
  4614  	}
  4615  
  4616  	ws := memdb.NewWatchSet()
  4617  	iter, err := state.AllocsByIDPrefix(ws, structs.DefaultNamespace, "aaaa")
  4618  	if err != nil {
  4619  		t.Fatalf("err: %v", err)
  4620  	}
  4621  
  4622  	gatherAllocs := func(iter memdb.ResultIterator) []*structs.Allocation {
  4623  		var allocs []*structs.Allocation
  4624  		for {
  4625  			raw := iter.Next()
  4626  			if raw == nil {
  4627  				break
  4628  			}
  4629  			allocs = append(allocs, raw.(*structs.Allocation))
  4630  		}
  4631  		return allocs
  4632  	}
  4633  
  4634  	out := gatherAllocs(iter)
  4635  	if len(out) != 5 {
  4636  		t.Fatalf("bad: expected five allocations, got: %#v", out)
  4637  	}
  4638  
  4639  	sort.Sort(AllocIDSort(allocs))
  4640  
  4641  	for index, alloc := range out {
  4642  		if ids[index] != alloc.ID {
  4643  			t.Fatalf("bad: got unexpected id: %s", alloc.ID)
  4644  		}
  4645  	}
  4646  
  4647  	iter, err = state.AllocsByIDPrefix(ws, structs.DefaultNamespace, "b-a7bfb")
  4648  	if err != nil {
  4649  		t.Fatalf("err: %v", err)
  4650  	}
  4651  
  4652  	out = gatherAllocs(iter)
  4653  	if len(out) != 0 {
  4654  		t.Fatalf("bad: unexpected zero allocations, got: %#v", out)
  4655  	}
  4656  
  4657  	if watchFired(ws) {
  4658  		t.Fatalf("bad")
  4659  	}
  4660  }
  4661  
  4662  func TestStateStore_Allocs(t *testing.T) {
  4663  	state := testStateStore(t)
  4664  	var allocs []*structs.Allocation
  4665  
  4666  	for i := 0; i < 10; i++ {
  4667  		alloc := mock.Alloc()
  4668  		allocs = append(allocs, alloc)
  4669  	}
  4670  	for i, alloc := range allocs {
  4671  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  4672  	}
  4673  
  4674  	err := state.UpsertAllocs(1000, allocs)
  4675  	if err != nil {
  4676  		t.Fatalf("err: %v", err)
  4677  	}
  4678  
  4679  	ws := memdb.NewWatchSet()
  4680  	iter, err := state.Allocs(ws)
  4681  	if err != nil {
  4682  		t.Fatalf("err: %v", err)
  4683  	}
  4684  
  4685  	var out []*structs.Allocation
  4686  	for {
  4687  		raw := iter.Next()
  4688  		if raw == nil {
  4689  			break
  4690  		}
  4691  		out = append(out, raw.(*structs.Allocation))
  4692  	}
  4693  
  4694  	sort.Sort(AllocIDSort(allocs))
  4695  	sort.Sort(AllocIDSort(out))
  4696  
  4697  	if !reflect.DeepEqual(allocs, out) {
  4698  		t.Fatalf("bad: %#v %#v", allocs, out)
  4699  	}
  4700  
  4701  	if watchFired(ws) {
  4702  		t.Fatalf("bad")
  4703  	}
  4704  }
  4705  
  4706  func TestStateStore_Allocs_PrevAlloc(t *testing.T) {
  4707  	state := testStateStore(t)
  4708  	var allocs []*structs.Allocation
  4709  
  4710  	require := require.New(t)
  4711  	for i := 0; i < 5; i++ {
  4712  		alloc := mock.Alloc()
  4713  		allocs = append(allocs, alloc)
  4714  	}
  4715  	for i, alloc := range allocs {
  4716  		state.UpsertJobSummary(uint64(900+i), mock.JobSummary(alloc.JobID))
  4717  	}
  4718  	// Set some previous alloc ids
  4719  	allocs[1].PreviousAllocation = allocs[0].ID
  4720  	allocs[2].PreviousAllocation = allocs[1].ID
  4721  
  4722  	err := state.UpsertAllocs(1000, allocs)
  4723  	require.Nil(err)
  4724  
  4725  	ws := memdb.NewWatchSet()
  4726  	iter, err := state.Allocs(ws)
  4727  	require.Nil(err)
  4728  
  4729  	var out []*structs.Allocation
  4730  	for {
  4731  		raw := iter.Next()
  4732  		if raw == nil {
  4733  			break
  4734  		}
  4735  		out = append(out, raw.(*structs.Allocation))
  4736  	}
  4737  
  4738  	// Set expected NextAllocation fields
  4739  	allocs[0].NextAllocation = allocs[1].ID
  4740  	allocs[1].NextAllocation = allocs[2].ID
  4741  
  4742  	sort.Sort(AllocIDSort(allocs))
  4743  	sort.Sort(AllocIDSort(out))
  4744  
  4745  	require.Equal(allocs, out)
  4746  	require.False(watchFired(ws))
  4747  
  4748  	// Insert another alloc, verify index of previous alloc also got updated
  4749  	alloc := mock.Alloc()
  4750  	alloc.PreviousAllocation = allocs[0].ID
  4751  	err = state.UpsertAllocs(1001, []*structs.Allocation{alloc})
  4752  	require.Nil(err)
  4753  	alloc0, err := state.AllocByID(nil, allocs[0].ID)
  4754  	require.Nil(err)
  4755  	require.Equal(alloc0.ModifyIndex, uint64(1001))
  4756  }
  4757  
  4758  func TestStateStore_RestoreAlloc(t *testing.T) {
  4759  	state := testStateStore(t)
  4760  	alloc := mock.Alloc()
  4761  
  4762  	restore, err := state.Restore()
  4763  	if err != nil {
  4764  		t.Fatalf("err: %v", err)
  4765  	}
  4766  
  4767  	err = restore.AllocRestore(alloc)
  4768  	if err != nil {
  4769  		t.Fatalf("err: %v", err)
  4770  	}
  4771  
  4772  	restore.Commit()
  4773  
  4774  	ws := memdb.NewWatchSet()
  4775  	out, err := state.AllocByID(ws, alloc.ID)
  4776  	if err != nil {
  4777  		t.Fatalf("err: %v", err)
  4778  	}
  4779  
  4780  	if !reflect.DeepEqual(out, alloc) {
  4781  		t.Fatalf("Bad: %#v %#v", out, alloc)
  4782  	}
  4783  
  4784  	if watchFired(ws) {
  4785  		t.Fatalf("bad")
  4786  	}
  4787  }
  4788  
  4789  func TestStateStore_RestoreAlloc_NoEphemeralDisk(t *testing.T) {
  4790  	state := testStateStore(t)
  4791  	alloc := mock.Alloc()
  4792  	alloc.Job.TaskGroups[0].EphemeralDisk = nil
  4793  	alloc.Job.TaskGroups[0].Tasks[0].Resources.DiskMB = 120
  4794  
  4795  	restore, err := state.Restore()
  4796  	if err != nil {
  4797  		t.Fatalf("err: %v", err)
  4798  	}
  4799  
  4800  	err = restore.AllocRestore(alloc)
  4801  	if err != nil {
  4802  		t.Fatalf("err: %v", err)
  4803  	}
  4804  
  4805  	restore.Commit()
  4806  
  4807  	ws := memdb.NewWatchSet()
  4808  	out, err := state.AllocByID(ws, alloc.ID)
  4809  	if err != nil {
  4810  		t.Fatalf("err: %v", err)
  4811  	}
  4812  
  4813  	expected := alloc.Copy()
  4814  	expected.Job.TaskGroups[0].EphemeralDisk = &structs.EphemeralDisk{SizeMB: 120}
  4815  	expected.Job.TaskGroups[0].Tasks[0].Resources.DiskMB = 0
  4816  
  4817  	if !reflect.DeepEqual(out, expected) {
  4818  		t.Fatalf("Bad: %#v %#v", out, expected)
  4819  	}
  4820  
  4821  	if watchFired(ws) {
  4822  		t.Fatalf("bad")
  4823  	}
  4824  }
  4825  
  4826  func TestStateStore_SetJobStatus_ForceStatus(t *testing.T) {
  4827  	state := testStateStore(t)
  4828  	txn := state.db.Txn(true)
  4829  
  4830  	// Create and insert a mock job.
  4831  	job := mock.Job()
  4832  	job.Status = ""
  4833  	job.ModifyIndex = 0
  4834  	if err := txn.Insert("jobs", job); err != nil {
  4835  		t.Fatalf("job insert failed: %v", err)
  4836  	}
  4837  
  4838  	exp := "foobar"
  4839  	index := uint64(1000)
  4840  	if err := state.setJobStatus(index, txn, job, false, exp); err != nil {
  4841  		t.Fatalf("setJobStatus() failed: %v", err)
  4842  	}
  4843  
  4844  	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
  4845  	if err != nil {
  4846  		t.Fatalf("job lookup failed: %v", err)
  4847  	}
  4848  	updated := i.(*structs.Job)
  4849  
  4850  	if updated.Status != exp {
  4851  		t.Fatalf("setJobStatus() set %v; expected %v", updated.Status, exp)
  4852  	}
  4853  
  4854  	if updated.ModifyIndex != index {
  4855  		t.Fatalf("setJobStatus() set %d; expected %d", updated.ModifyIndex, index)
  4856  	}
  4857  }
  4858  
  4859  func TestStateStore_SetJobStatus_NoOp(t *testing.T) {
  4860  	state := testStateStore(t)
  4861  	txn := state.db.Txn(true)
  4862  
  4863  	// Create and insert a mock job that should be pending.
  4864  	job := mock.Job()
  4865  	job.Status = structs.JobStatusPending
  4866  	job.ModifyIndex = 10
  4867  	if err := txn.Insert("jobs", job); err != nil {
  4868  		t.Fatalf("job insert failed: %v", err)
  4869  	}
  4870  
  4871  	index := uint64(1000)
  4872  	if err := state.setJobStatus(index, txn, job, false, ""); err != nil {
  4873  		t.Fatalf("setJobStatus() failed: %v", err)
  4874  	}
  4875  
  4876  	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
  4877  	if err != nil {
  4878  		t.Fatalf("job lookup failed: %v", err)
  4879  	}
  4880  	updated := i.(*structs.Job)
  4881  
  4882  	if updated.ModifyIndex == index {
  4883  		t.Fatalf("setJobStatus() should have been a no-op")
  4884  	}
  4885  }
  4886  
  4887  func TestStateStore_SetJobStatus(t *testing.T) {
  4888  	state := testStateStore(t)
  4889  	txn := state.db.Txn(true)
  4890  
  4891  	// Create and insert a mock job that should be pending but has an incorrect
  4892  	// status.
  4893  	job := mock.Job()
  4894  	job.Status = "foobar"
  4895  	job.ModifyIndex = 10
  4896  	if err := txn.Insert("jobs", job); err != nil {
  4897  		t.Fatalf("job insert failed: %v", err)
  4898  	}
  4899  
  4900  	index := uint64(1000)
  4901  	if err := state.setJobStatus(index, txn, job, false, ""); err != nil {
  4902  		t.Fatalf("setJobStatus() failed: %v", err)
  4903  	}
  4904  
  4905  	i, err := txn.First("jobs", "id", job.Namespace, job.ID)
  4906  	if err != nil {
  4907  		t.Fatalf("job lookup failed: %v", err)
  4908  	}
  4909  	updated := i.(*structs.Job)
  4910  
  4911  	if updated.Status != structs.JobStatusPending {
  4912  		t.Fatalf("setJobStatus() set %v; expected %v", updated.Status, structs.JobStatusPending)
  4913  	}
  4914  
  4915  	if updated.ModifyIndex != index {
  4916  		t.Fatalf("setJobStatus() set %d; expected %d", updated.ModifyIndex, index)
  4917  	}
  4918  }
  4919  
  4920  func TestStateStore_GetJobStatus_NoEvalsOrAllocs(t *testing.T) {
  4921  	job := mock.Job()
  4922  	state := testStateStore(t)
  4923  	txn := state.db.Txn(false)
  4924  	status, err := state.getJobStatus(txn, job, false)
  4925  	if err != nil {
  4926  		t.Fatalf("getJobStatus() failed: %v", err)
  4927  	}
  4928  
  4929  	if status != structs.JobStatusPending {
  4930  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusPending)
  4931  	}
  4932  }
  4933  
  4934  func TestStateStore_GetJobStatus_NoEvalsOrAllocs_Periodic(t *testing.T) {
  4935  	job := mock.PeriodicJob()
  4936  	state := testStateStore(t)
  4937  	txn := state.db.Txn(false)
  4938  	status, err := state.getJobStatus(txn, job, false)
  4939  	if err != nil {
  4940  		t.Fatalf("getJobStatus() failed: %v", err)
  4941  	}
  4942  
  4943  	if status != structs.JobStatusRunning {
  4944  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  4945  	}
  4946  }
  4947  
  4948  func TestStateStore_GetJobStatus_NoEvalsOrAllocs_EvalDelete(t *testing.T) {
  4949  	job := mock.Job()
  4950  	state := testStateStore(t)
  4951  	txn := state.db.Txn(false)
  4952  	status, err := state.getJobStatus(txn, job, true)
  4953  	if err != nil {
  4954  		t.Fatalf("getJobStatus() failed: %v", err)
  4955  	}
  4956  
  4957  	if status != structs.JobStatusDead {
  4958  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  4959  	}
  4960  }
  4961  
  4962  func TestStateStore_GetJobStatus_DeadEvalsAndAllocs(t *testing.T) {
  4963  	state := testStateStore(t)
  4964  	job := mock.Job()
  4965  
  4966  	// Create a mock alloc that is dead.
  4967  	alloc := mock.Alloc()
  4968  	alloc.JobID = job.ID
  4969  	alloc.DesiredStatus = structs.AllocDesiredStatusStop
  4970  	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
  4971  	if err := state.UpsertAllocs(1000, []*structs.Allocation{alloc}); err != nil {
  4972  		t.Fatalf("err: %v", err)
  4973  	}
  4974  
  4975  	// Create a mock eval that is complete
  4976  	eval := mock.Eval()
  4977  	eval.JobID = job.ID
  4978  	eval.Status = structs.EvalStatusComplete
  4979  	if err := state.UpsertEvals(1001, []*structs.Evaluation{eval}); err != nil {
  4980  		t.Fatalf("err: %v", err)
  4981  	}
  4982  
  4983  	txn := state.db.Txn(false)
  4984  	status, err := state.getJobStatus(txn, job, false)
  4985  	if err != nil {
  4986  		t.Fatalf("getJobStatus() failed: %v", err)
  4987  	}
  4988  
  4989  	if status != structs.JobStatusDead {
  4990  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  4991  	}
  4992  }
  4993  
  4994  func TestStateStore_GetJobStatus_RunningAlloc(t *testing.T) {
  4995  	state := testStateStore(t)
  4996  	job := mock.Job()
  4997  
  4998  	// Create a mock alloc that is running.
  4999  	alloc := mock.Alloc()
  5000  	alloc.JobID = job.ID
  5001  	alloc.DesiredStatus = structs.AllocDesiredStatusRun
  5002  	state.UpsertJobSummary(999, mock.JobSummary(alloc.JobID))
  5003  	if err := state.UpsertAllocs(1000, []*structs.Allocation{alloc}); err != nil {
  5004  		t.Fatalf("err: %v", err)
  5005  	}
  5006  
  5007  	txn := state.db.Txn(false)
  5008  	status, err := state.getJobStatus(txn, job, true)
  5009  	if err != nil {
  5010  		t.Fatalf("getJobStatus() failed: %v", err)
  5011  	}
  5012  
  5013  	if status != structs.JobStatusRunning {
  5014  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  5015  	}
  5016  }
  5017  
  5018  func TestStateStore_GetJobStatus_PeriodicJob(t *testing.T) {
  5019  	state := testStateStore(t)
  5020  	job := mock.PeriodicJob()
  5021  
  5022  	txn := state.db.Txn(false)
  5023  	status, err := state.getJobStatus(txn, job, false)
  5024  	if err != nil {
  5025  		t.Fatalf("getJobStatus() failed: %v", err)
  5026  	}
  5027  
  5028  	if status != structs.JobStatusRunning {
  5029  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  5030  	}
  5031  
  5032  	// Mark it as stopped
  5033  	job.Stop = true
  5034  	status, err = state.getJobStatus(txn, job, false)
  5035  	if err != nil {
  5036  		t.Fatalf("getJobStatus() failed: %v", err)
  5037  	}
  5038  
  5039  	if status != structs.JobStatusDead {
  5040  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  5041  	}
  5042  }
  5043  
  5044  func TestStateStore_GetJobStatus_ParameterizedJob(t *testing.T) {
  5045  	state := testStateStore(t)
  5046  	job := mock.Job()
  5047  	job.ParameterizedJob = &structs.ParameterizedJobConfig{}
  5048  
  5049  	txn := state.db.Txn(false)
  5050  	status, err := state.getJobStatus(txn, job, false)
  5051  	if err != nil {
  5052  		t.Fatalf("getJobStatus() failed: %v", err)
  5053  	}
  5054  
  5055  	if status != structs.JobStatusRunning {
  5056  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusRunning)
  5057  	}
  5058  
  5059  	// Mark it as stopped
  5060  	job.Stop = true
  5061  	status, err = state.getJobStatus(txn, job, false)
  5062  	if err != nil {
  5063  		t.Fatalf("getJobStatus() failed: %v", err)
  5064  	}
  5065  
  5066  	if status != structs.JobStatusDead {
  5067  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusDead)
  5068  	}
  5069  }
  5070  
  5071  func TestStateStore_SetJobStatus_PendingEval(t *testing.T) {
  5072  	state := testStateStore(t)
  5073  	job := mock.Job()
  5074  
  5075  	// Create a mock eval that is pending.
  5076  	eval := mock.Eval()
  5077  	eval.JobID = job.ID
  5078  	eval.Status = structs.EvalStatusPending
  5079  	if err := state.UpsertEvals(1000, []*structs.Evaluation{eval}); err != nil {
  5080  		t.Fatalf("err: %v", err)
  5081  	}
  5082  
  5083  	txn := state.db.Txn(false)
  5084  	status, err := state.getJobStatus(txn, job, true)
  5085  	if err != nil {
  5086  		t.Fatalf("getJobStatus() failed: %v", err)
  5087  	}
  5088  
  5089  	if status != structs.JobStatusPending {
  5090  		t.Fatalf("getJobStatus() returned %v; expected %v", status, structs.JobStatusPending)
  5091  	}
  5092  }
  5093  
  5094  // TestStateStore_SetJobStatus_SystemJob asserts that system jobs are still
  5095  // considered running until explicitly stopped.
  5096  func TestStateStore_SetJobStatus_SystemJob(t *testing.T) {
  5097  	state := testStateStore(t)
  5098  	job := mock.SystemJob()
  5099  
  5100  	// Create a mock eval that is pending.
  5101  	eval := mock.Eval()
  5102  	eval.JobID = job.ID
  5103  	eval.Type = job.Type
  5104  	eval.Status = structs.EvalStatusComplete
  5105  	if err := state.UpsertEvals(1000, []*structs.Evaluation{eval}); err != nil {
  5106  		t.Fatalf("err: %v", err)
  5107  	}
  5108  
  5109  	txn := state.db.Txn(false)
  5110  	status, err := state.getJobStatus(txn, job, true)
  5111  	if err != nil {
  5112  		t.Fatalf("getJobStatus() failed: %v", err)
  5113  	}
  5114  
  5115  	if expected := structs.JobStatusRunning; status != expected {
  5116  		t.Fatalf("getJobStatus() returned %v; expected %v", status, expected)
  5117  	}
  5118  
  5119  	// Stop the job
  5120  	job.Stop = true
  5121  	status, err = state.getJobStatus(txn, job, true)
  5122  	if err != nil {
  5123  		t.Fatalf("getJobStatus() failed: %v", err)
  5124  	}
  5125  
  5126  	if expected := structs.JobStatusDead; status != expected {
  5127  		t.Fatalf("getJobStatus() returned %v; expected %v", status, expected)
  5128  	}
  5129  }
  5130  
  5131  func TestStateJobSummary_UpdateJobCount(t *testing.T) {
  5132  	state := testStateStore(t)
  5133  	alloc := mock.Alloc()
  5134  	job := alloc.Job
  5135  	job.TaskGroups[0].Count = 3
  5136  
  5137  	// Create watchsets so we can test that upsert fires the watch
  5138  	ws := memdb.NewWatchSet()
  5139  	if _, err := state.JobSummaryByID(ws, job.Namespace, job.ID); err != nil {
  5140  		t.Fatalf("bad: %v", err)
  5141  	}
  5142  
  5143  	if err := state.UpsertJob(1000, job); err != nil {
  5144  		t.Fatalf("err: %v", err)
  5145  	}
  5146  
  5147  	if err := state.UpsertAllocs(1001, []*structs.Allocation{alloc}); 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, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
  5157  	expectedSummary := structs.JobSummary{
  5158  		JobID:     job.ID,
  5159  		Namespace: job.Namespace,
  5160  		Summary: map[string]structs.TaskGroupSummary{
  5161  			"web": {
  5162  				Starting: 1,
  5163  			},
  5164  		},
  5165  		Children:    new(structs.JobChildrenSummary),
  5166  		CreateIndex: 1000,
  5167  		ModifyIndex: 1001,
  5168  	}
  5169  	if !reflect.DeepEqual(summary, &expectedSummary) {
  5170  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  5171  	}
  5172  
  5173  	// Create watchsets so we can test that upsert fires the watch
  5174  	ws2 := memdb.NewWatchSet()
  5175  	if _, err := state.JobSummaryByID(ws2, job.Namespace, job.ID); err != nil {
  5176  		t.Fatalf("bad: %v", err)
  5177  	}
  5178  
  5179  	alloc2 := mock.Alloc()
  5180  	alloc2.Job = job
  5181  	alloc2.JobID = job.ID
  5182  
  5183  	alloc3 := mock.Alloc()
  5184  	alloc3.Job = job
  5185  	alloc3.JobID = job.ID
  5186  
  5187  	if err := state.UpsertAllocs(1002, []*structs.Allocation{alloc2, alloc3}); err != nil {
  5188  		t.Fatalf("err: %v", err)
  5189  	}
  5190  
  5191  	if !watchFired(ws2) {
  5192  		t.Fatalf("bad")
  5193  	}
  5194  
  5195  	outA, _ := state.AllocByID(ws, alloc3.ID)
  5196  
  5197  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  5198  	expectedSummary = structs.JobSummary{
  5199  		JobID:     job.ID,
  5200  		Namespace: job.Namespace,
  5201  		Summary: map[string]structs.TaskGroupSummary{
  5202  			"web": {
  5203  				Starting: 3,
  5204  			},
  5205  		},
  5206  		Children:    new(structs.JobChildrenSummary),
  5207  		CreateIndex: job.CreateIndex,
  5208  		ModifyIndex: outA.ModifyIndex,
  5209  	}
  5210  	if !reflect.DeepEqual(summary, &expectedSummary) {
  5211  		t.Fatalf("expected summary: %v, actual: %v", expectedSummary, summary)
  5212  	}
  5213  
  5214  	// Create watchsets so we can test that upsert fires the watch
  5215  	ws3 := memdb.NewWatchSet()
  5216  	if _, err := state.JobSummaryByID(ws3, job.Namespace, job.ID); err != nil {
  5217  		t.Fatalf("bad: %v", err)
  5218  	}
  5219  
  5220  	alloc4 := mock.Alloc()
  5221  	alloc4.ID = alloc2.ID
  5222  	alloc4.Job = alloc2.Job
  5223  	alloc4.JobID = alloc2.JobID
  5224  	alloc4.ClientStatus = structs.AllocClientStatusComplete
  5225  
  5226  	alloc5 := mock.Alloc()
  5227  	alloc5.ID = alloc3.ID
  5228  	alloc5.Job = alloc3.Job
  5229  	alloc5.JobID = alloc3.JobID
  5230  	alloc5.ClientStatus = structs.AllocClientStatusComplete
  5231  
  5232  	if err := state.UpdateAllocsFromClient(1004, []*structs.Allocation{alloc4, alloc5}); err != nil {
  5233  		t.Fatalf("err: %v", err)
  5234  	}
  5235  
  5236  	if !watchFired(ws2) {
  5237  		t.Fatalf("bad")
  5238  	}
  5239  
  5240  	outA, _ = state.AllocByID(ws, alloc5.ID)
  5241  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  5242  	expectedSummary = structs.JobSummary{
  5243  		JobID:     job.ID,
  5244  		Namespace: job.Namespace,
  5245  		Summary: map[string]structs.TaskGroupSummary{
  5246  			"web": {
  5247  				Complete: 2,
  5248  				Starting: 1,
  5249  			},
  5250  		},
  5251  		Children:    new(structs.JobChildrenSummary),
  5252  		CreateIndex: job.CreateIndex,
  5253  		ModifyIndex: outA.ModifyIndex,
  5254  	}
  5255  	if !reflect.DeepEqual(summary, &expectedSummary) {
  5256  		t.Fatalf("expected: %v, actual: %v", expectedSummary, summary)
  5257  	}
  5258  }
  5259  
  5260  func TestJobSummary_UpdateClientStatus(t *testing.T) {
  5261  	state := testStateStore(t)
  5262  	alloc := mock.Alloc()
  5263  	job := alloc.Job
  5264  	job.TaskGroups[0].Count = 3
  5265  
  5266  	alloc2 := mock.Alloc()
  5267  	alloc2.Job = job
  5268  	alloc2.JobID = job.ID
  5269  
  5270  	alloc3 := mock.Alloc()
  5271  	alloc3.Job = job
  5272  	alloc3.JobID = job.ID
  5273  
  5274  	err := state.UpsertJob(1000, job)
  5275  	if err != nil {
  5276  		t.Fatalf("err: %v", err)
  5277  	}
  5278  
  5279  	if err := state.UpsertAllocs(1001, []*structs.Allocation{alloc, alloc2, alloc3}); err != nil {
  5280  		t.Fatalf("err: %v", err)
  5281  	}
  5282  
  5283  	ws := memdb.NewWatchSet()
  5284  	summary, _ := state.JobSummaryByID(ws, job.Namespace, job.ID)
  5285  	if summary.Summary["web"].Starting != 3 {
  5286  		t.Fatalf("bad job summary: %v", summary)
  5287  	}
  5288  
  5289  	alloc4 := mock.Alloc()
  5290  	alloc4.ID = alloc2.ID
  5291  	alloc4.Job = alloc2.Job
  5292  	alloc4.JobID = alloc2.JobID
  5293  	alloc4.ClientStatus = structs.AllocClientStatusComplete
  5294  
  5295  	alloc5 := mock.Alloc()
  5296  	alloc5.ID = alloc3.ID
  5297  	alloc5.Job = alloc3.Job
  5298  	alloc5.JobID = alloc3.JobID
  5299  	alloc5.ClientStatus = structs.AllocClientStatusFailed
  5300  
  5301  	alloc6 := mock.Alloc()
  5302  	alloc6.ID = alloc.ID
  5303  	alloc6.Job = alloc.Job
  5304  	alloc6.JobID = alloc.JobID
  5305  	alloc6.ClientStatus = structs.AllocClientStatusRunning
  5306  
  5307  	if err := state.UpdateAllocsFromClient(1002, []*structs.Allocation{alloc4, alloc5, alloc6}); err != nil {
  5308  		t.Fatalf("err: %v", err)
  5309  	}
  5310  
  5311  	if !watchFired(ws) {
  5312  		t.Fatalf("bad")
  5313  	}
  5314  
  5315  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  5316  	if summary.Summary["web"].Running != 1 || summary.Summary["web"].Failed != 1 || summary.Summary["web"].Complete != 1 {
  5317  		t.Fatalf("bad job summary: %v", summary)
  5318  	}
  5319  
  5320  	alloc7 := mock.Alloc()
  5321  	alloc7.Job = alloc.Job
  5322  	alloc7.JobID = alloc.JobID
  5323  
  5324  	if err := state.UpsertAllocs(1003, []*structs.Allocation{alloc7}); err != nil {
  5325  		t.Fatalf("err: %v", err)
  5326  	}
  5327  	summary, _ = state.JobSummaryByID(ws, job.Namespace, job.ID)
  5328  	if summary.Summary["web"].Starting != 1 || summary.Summary["web"].Running != 1 || summary.Summary["web"].Failed != 1 || summary.Summary["web"].Complete != 1 {
  5329  		t.Fatalf("bad job summary: %v", summary)
  5330  	}
  5331  }
  5332  
  5333  // Test that nonexistent deployment can't be updated
  5334  func TestStateStore_UpsertDeploymentStatusUpdate_Nonexistent(t *testing.T) {
  5335  	state := testStateStore(t)
  5336  
  5337  	// Update the nonexistent deployment
  5338  	req := &structs.DeploymentStatusUpdateRequest{
  5339  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  5340  			DeploymentID: uuid.Generate(),
  5341  			Status:       structs.DeploymentStatusRunning,
  5342  		},
  5343  	}
  5344  	err := state.UpdateDeploymentStatus(2, req)
  5345  	if err == nil || !strings.Contains(err.Error(), "does not exist") {
  5346  		t.Fatalf("expected error updating the status because the deployment doesn't exist")
  5347  	}
  5348  }
  5349  
  5350  // Test that terminal deployment can't be updated
  5351  func TestStateStore_UpsertDeploymentStatusUpdate_Terminal(t *testing.T) {
  5352  	state := testStateStore(t)
  5353  
  5354  	// Insert a terminal deployment
  5355  	d := mock.Deployment()
  5356  	d.Status = structs.DeploymentStatusFailed
  5357  
  5358  	if err := state.UpsertDeployment(1, d); err != nil {
  5359  		t.Fatalf("bad: %v", err)
  5360  	}
  5361  
  5362  	// Update the deployment
  5363  	req := &structs.DeploymentStatusUpdateRequest{
  5364  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  5365  			DeploymentID: d.ID,
  5366  			Status:       structs.DeploymentStatusRunning,
  5367  		},
  5368  	}
  5369  	err := state.UpdateDeploymentStatus(2, req)
  5370  	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
  5371  		t.Fatalf("expected error updating the status because the deployment is terminal")
  5372  	}
  5373  }
  5374  
  5375  // Test that a non terminal deployment is updated and that a job and eval are
  5376  // created.
  5377  func TestStateStore_UpsertDeploymentStatusUpdate_NonTerminal(t *testing.T) {
  5378  	state := testStateStore(t)
  5379  
  5380  	// Insert a deployment
  5381  	d := mock.Deployment()
  5382  	if err := state.UpsertDeployment(1, d); err != nil {
  5383  		t.Fatalf("bad: %v", err)
  5384  	}
  5385  
  5386  	// Create an eval and a job
  5387  	e := mock.Eval()
  5388  	j := mock.Job()
  5389  
  5390  	// Update the deployment
  5391  	status, desc := structs.DeploymentStatusFailed, "foo"
  5392  	req := &structs.DeploymentStatusUpdateRequest{
  5393  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  5394  			DeploymentID:      d.ID,
  5395  			Status:            status,
  5396  			StatusDescription: desc,
  5397  		},
  5398  		Job:  j,
  5399  		Eval: e,
  5400  	}
  5401  	err := state.UpdateDeploymentStatus(2, req)
  5402  	if err != nil {
  5403  		t.Fatalf("bad: %v", err)
  5404  	}
  5405  
  5406  	// Check that the status was updated properly
  5407  	ws := memdb.NewWatchSet()
  5408  	dout, err := state.DeploymentByID(ws, d.ID)
  5409  	if err != nil {
  5410  		t.Fatalf("bad: %v", err)
  5411  	}
  5412  	if dout.Status != status || dout.StatusDescription != desc {
  5413  		t.Fatalf("bad: %#v", dout)
  5414  	}
  5415  
  5416  	// Check that the evaluation was created
  5417  	eout, _ := state.EvalByID(ws, e.ID)
  5418  	if err != nil {
  5419  		t.Fatalf("bad: %v", err)
  5420  	}
  5421  	if eout == nil {
  5422  		t.Fatalf("bad: %#v", eout)
  5423  	}
  5424  
  5425  	// Check that the job was created
  5426  	jout, _ := state.JobByID(ws, j.Namespace, j.ID)
  5427  	if err != nil {
  5428  		t.Fatalf("bad: %v", err)
  5429  	}
  5430  	if jout == nil {
  5431  		t.Fatalf("bad: %#v", jout)
  5432  	}
  5433  }
  5434  
  5435  // Test that when a deployment is updated to successful the job is updated to
  5436  // stable
  5437  func TestStateStore_UpsertDeploymentStatusUpdate_Successful(t *testing.T) {
  5438  	state := testStateStore(t)
  5439  
  5440  	// Insert a job
  5441  	job := mock.Job()
  5442  	if err := state.UpsertJob(1, job); err != nil {
  5443  		t.Fatalf("bad: %v", err)
  5444  	}
  5445  
  5446  	// Insert a deployment
  5447  	d := structs.NewDeployment(job)
  5448  	if err := state.UpsertDeployment(2, d); err != nil {
  5449  		t.Fatalf("bad: %v", err)
  5450  	}
  5451  
  5452  	// Update the deployment
  5453  	req := &structs.DeploymentStatusUpdateRequest{
  5454  		DeploymentUpdate: &structs.DeploymentStatusUpdate{
  5455  			DeploymentID:      d.ID,
  5456  			Status:            structs.DeploymentStatusSuccessful,
  5457  			StatusDescription: structs.DeploymentStatusDescriptionSuccessful,
  5458  		},
  5459  	}
  5460  	err := state.UpdateDeploymentStatus(3, req)
  5461  	if err != nil {
  5462  		t.Fatalf("bad: %v", err)
  5463  	}
  5464  
  5465  	// Check that the status was updated properly
  5466  	ws := memdb.NewWatchSet()
  5467  	dout, err := state.DeploymentByID(ws, d.ID)
  5468  	if err != nil {
  5469  		t.Fatalf("bad: %v", err)
  5470  	}
  5471  	if dout.Status != structs.DeploymentStatusSuccessful ||
  5472  		dout.StatusDescription != structs.DeploymentStatusDescriptionSuccessful {
  5473  		t.Fatalf("bad: %#v", dout)
  5474  	}
  5475  
  5476  	// Check that the job was created
  5477  	jout, _ := state.JobByID(ws, job.Namespace, job.ID)
  5478  	if err != nil {
  5479  		t.Fatalf("bad: %v", err)
  5480  	}
  5481  	if jout == nil {
  5482  		t.Fatalf("bad: %#v", jout)
  5483  	}
  5484  	if !jout.Stable {
  5485  		t.Fatalf("job not marked stable %#v", jout)
  5486  	}
  5487  	if jout.Version != d.JobVersion {
  5488  		t.Fatalf("job version changed; got %d; want %d", jout.Version, d.JobVersion)
  5489  	}
  5490  }
  5491  
  5492  func TestStateStore_UpdateJobStability(t *testing.T) {
  5493  	state := testStateStore(t)
  5494  
  5495  	// Insert a job twice to get two versions
  5496  	job := mock.Job()
  5497  	if err := state.UpsertJob(1, job); err != nil {
  5498  		t.Fatalf("bad: %v", err)
  5499  	}
  5500  
  5501  	if err := state.UpsertJob(2, job); err != nil {
  5502  		t.Fatalf("bad: %v", err)
  5503  	}
  5504  
  5505  	// Update the stability to true
  5506  	err := state.UpdateJobStability(3, job.Namespace, job.ID, 0, true)
  5507  	if err != nil {
  5508  		t.Fatalf("bad: %v", err)
  5509  	}
  5510  
  5511  	// Check that the job was updated properly
  5512  	ws := memdb.NewWatchSet()
  5513  	jout, _ := state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
  5514  	if err != nil {
  5515  		t.Fatalf("bad: %v", err)
  5516  	}
  5517  	if jout == nil {
  5518  		t.Fatalf("bad: %#v", jout)
  5519  	}
  5520  	if !jout.Stable {
  5521  		t.Fatalf("job not marked stable %#v", jout)
  5522  	}
  5523  
  5524  	// Update the stability to false
  5525  	err = state.UpdateJobStability(3, job.Namespace, job.ID, 0, false)
  5526  	if err != nil {
  5527  		t.Fatalf("bad: %v", err)
  5528  	}
  5529  
  5530  	// Check that the job was updated properly
  5531  	jout, _ = state.JobByIDAndVersion(ws, job.Namespace, job.ID, 0)
  5532  	if err != nil {
  5533  		t.Fatalf("bad: %v", err)
  5534  	}
  5535  	if jout == nil {
  5536  		t.Fatalf("bad: %#v", jout)
  5537  	}
  5538  	if jout.Stable {
  5539  		t.Fatalf("job marked stable %#v", jout)
  5540  	}
  5541  }
  5542  
  5543  // Test that nonexistent deployment can't be promoted
  5544  func TestStateStore_UpsertDeploymentPromotion_Nonexistent(t *testing.T) {
  5545  	state := testStateStore(t)
  5546  
  5547  	// Promote the nonexistent deployment
  5548  	req := &structs.ApplyDeploymentPromoteRequest{
  5549  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  5550  			DeploymentID: uuid.Generate(),
  5551  			All:          true,
  5552  		},
  5553  	}
  5554  	err := state.UpdateDeploymentPromotion(2, req)
  5555  	if err == nil || !strings.Contains(err.Error(), "does not exist") {
  5556  		t.Fatalf("expected error promoting because the deployment doesn't exist")
  5557  	}
  5558  }
  5559  
  5560  // Test that terminal deployment can't be updated
  5561  func TestStateStore_UpsertDeploymentPromotion_Terminal(t *testing.T) {
  5562  	state := testStateStore(t)
  5563  
  5564  	// Insert a terminal deployment
  5565  	d := mock.Deployment()
  5566  	d.Status = structs.DeploymentStatusFailed
  5567  
  5568  	if err := state.UpsertDeployment(1, d); err != nil {
  5569  		t.Fatalf("bad: %v", err)
  5570  	}
  5571  
  5572  	// Promote the deployment
  5573  	req := &structs.ApplyDeploymentPromoteRequest{
  5574  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  5575  			DeploymentID: d.ID,
  5576  			All:          true,
  5577  		},
  5578  	}
  5579  	err := state.UpdateDeploymentPromotion(2, req)
  5580  	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
  5581  		t.Fatalf("expected error updating the status because the deployment is terminal: %v", err)
  5582  	}
  5583  }
  5584  
  5585  // Test promoting unhealthy canaries in a deployment.
  5586  func TestStateStore_UpsertDeploymentPromotion_Unhealthy(t *testing.T) {
  5587  	state := testStateStore(t)
  5588  	require := require.New(t)
  5589  
  5590  	// Create a job
  5591  	j := mock.Job()
  5592  	require.Nil(state.UpsertJob(1, j))
  5593  
  5594  	// Create a deployment
  5595  	d := mock.Deployment()
  5596  	d.JobID = j.ID
  5597  	d.TaskGroups["web"].DesiredCanaries = 2
  5598  	require.Nil(state.UpsertDeployment(2, d))
  5599  
  5600  	// Create a set of allocations
  5601  	c1 := mock.Alloc()
  5602  	c1.JobID = j.ID
  5603  	c1.DeploymentID = d.ID
  5604  	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
  5605  	c2 := mock.Alloc()
  5606  	c2.JobID = j.ID
  5607  	c2.DeploymentID = d.ID
  5608  	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
  5609  
  5610  	require.Nil(state.UpsertAllocs(3, []*structs.Allocation{c1, c2}))
  5611  
  5612  	// Promote the canaries
  5613  	req := &structs.ApplyDeploymentPromoteRequest{
  5614  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  5615  			DeploymentID: d.ID,
  5616  			All:          true,
  5617  		},
  5618  	}
  5619  	err := state.UpdateDeploymentPromotion(4, req)
  5620  	require.NotNil(err)
  5621  	require.Contains(err.Error(), `Task group "web" has 0/2 healthy allocations`)
  5622  }
  5623  
  5624  // Test promoting a deployment with no canaries
  5625  func TestStateStore_UpsertDeploymentPromotion_NoCanaries(t *testing.T) {
  5626  	state := testStateStore(t)
  5627  	require := require.New(t)
  5628  
  5629  	// Create a job
  5630  	j := mock.Job()
  5631  	require.Nil(state.UpsertJob(1, j))
  5632  
  5633  	// Create a deployment
  5634  	d := mock.Deployment()
  5635  	d.TaskGroups["web"].DesiredCanaries = 2
  5636  	d.JobID = j.ID
  5637  	require.Nil(state.UpsertDeployment(2, d))
  5638  
  5639  	// Promote the canaries
  5640  	req := &structs.ApplyDeploymentPromoteRequest{
  5641  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  5642  			DeploymentID: d.ID,
  5643  			All:          true,
  5644  		},
  5645  	}
  5646  	err := state.UpdateDeploymentPromotion(4, req)
  5647  	require.NotNil(err)
  5648  	require.Contains(err.Error(), `Task group "web" has 0/2 healthy allocations`)
  5649  }
  5650  
  5651  // Test promoting all canaries in a deployment.
  5652  func TestStateStore_UpsertDeploymentPromotion_All(t *testing.T) {
  5653  	state := testStateStore(t)
  5654  
  5655  	// Create a job with two task groups
  5656  	j := mock.Job()
  5657  	tg1 := j.TaskGroups[0]
  5658  	tg2 := tg1.Copy()
  5659  	tg2.Name = "foo"
  5660  	j.TaskGroups = append(j.TaskGroups, tg2)
  5661  	if err := state.UpsertJob(1, j); err != nil {
  5662  		t.Fatalf("bad: %v", err)
  5663  	}
  5664  
  5665  	// Create a deployment
  5666  	d := mock.Deployment()
  5667  	d.StatusDescription = structs.DeploymentStatusDescriptionRunningNeedsPromotion
  5668  	d.JobID = j.ID
  5669  	d.TaskGroups = map[string]*structs.DeploymentState{
  5670  		"web": {
  5671  			DesiredTotal:    10,
  5672  			DesiredCanaries: 1,
  5673  		},
  5674  		"foo": {
  5675  			DesiredTotal:    10,
  5676  			DesiredCanaries: 1,
  5677  		},
  5678  	}
  5679  	if err := state.UpsertDeployment(2, d); err != nil {
  5680  		t.Fatalf("bad: %v", err)
  5681  	}
  5682  
  5683  	// Create a set of allocations
  5684  	c1 := mock.Alloc()
  5685  	c1.JobID = j.ID
  5686  	c1.DeploymentID = d.ID
  5687  	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
  5688  	c1.DeploymentStatus = &structs.AllocDeploymentStatus{
  5689  		Healthy: helper.BoolToPtr(true),
  5690  	}
  5691  	c2 := mock.Alloc()
  5692  	c2.JobID = j.ID
  5693  	c2.DeploymentID = d.ID
  5694  	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
  5695  	c2.TaskGroup = tg2.Name
  5696  	c2.DeploymentStatus = &structs.AllocDeploymentStatus{
  5697  		Healthy: helper.BoolToPtr(true),
  5698  	}
  5699  
  5700  	if err := state.UpsertAllocs(3, []*structs.Allocation{c1, c2}); err != nil {
  5701  		t.Fatalf("err: %v", err)
  5702  	}
  5703  
  5704  	// Create an eval
  5705  	e := mock.Eval()
  5706  
  5707  	// Promote the canaries
  5708  	req := &structs.ApplyDeploymentPromoteRequest{
  5709  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  5710  			DeploymentID: d.ID,
  5711  			All:          true,
  5712  		},
  5713  		Eval: e,
  5714  	}
  5715  	err := state.UpdateDeploymentPromotion(4, req)
  5716  	if err != nil {
  5717  		t.Fatalf("bad: %v", err)
  5718  	}
  5719  
  5720  	// Check that the status per task group was updated properly
  5721  	ws := memdb.NewWatchSet()
  5722  	dout, err := state.DeploymentByID(ws, d.ID)
  5723  	if err != nil {
  5724  		t.Fatalf("bad: %v", err)
  5725  	}
  5726  	if dout.StatusDescription != structs.DeploymentStatusDescriptionRunning {
  5727  		t.Fatalf("status description not updated: got %v; want %v", dout.StatusDescription, structs.DeploymentStatusDescriptionRunning)
  5728  	}
  5729  	if len(dout.TaskGroups) != 2 {
  5730  		t.Fatalf("bad: %#v", dout.TaskGroups)
  5731  	}
  5732  	for tg, state := range dout.TaskGroups {
  5733  		if !state.Promoted {
  5734  			t.Fatalf("bad: group %q not promoted %#v", tg, state)
  5735  		}
  5736  	}
  5737  
  5738  	// Check that the evaluation was created
  5739  	eout, _ := state.EvalByID(ws, e.ID)
  5740  	if err != nil {
  5741  		t.Fatalf("bad: %v", err)
  5742  	}
  5743  	if eout == nil {
  5744  		t.Fatalf("bad: %#v", eout)
  5745  	}
  5746  }
  5747  
  5748  // Test promoting a subset of canaries in a deployment.
  5749  func TestStateStore_UpsertDeploymentPromotion_Subset(t *testing.T) {
  5750  	state := testStateStore(t)
  5751  	require := require.New(t)
  5752  
  5753  	// Create a job with two task groups
  5754  	j := mock.Job()
  5755  	tg1 := j.TaskGroups[0]
  5756  	tg2 := tg1.Copy()
  5757  	tg2.Name = "foo"
  5758  	j.TaskGroups = append(j.TaskGroups, tg2)
  5759  	require.Nil(state.UpsertJob(1, j))
  5760  
  5761  	// Create a deployment
  5762  	d := mock.Deployment()
  5763  	d.JobID = j.ID
  5764  	d.TaskGroups = map[string]*structs.DeploymentState{
  5765  		"web": {
  5766  			DesiredTotal:    10,
  5767  			DesiredCanaries: 1,
  5768  		},
  5769  		"foo": {
  5770  			DesiredTotal:    10,
  5771  			DesiredCanaries: 1,
  5772  		},
  5773  	}
  5774  	require.Nil(state.UpsertDeployment(2, d))
  5775  
  5776  	// Create a set of allocations for both groups, including an unhealthy one
  5777  	c1 := mock.Alloc()
  5778  	c1.JobID = j.ID
  5779  	c1.DeploymentID = d.ID
  5780  	d.TaskGroups[c1.TaskGroup].PlacedCanaries = append(d.TaskGroups[c1.TaskGroup].PlacedCanaries, c1.ID)
  5781  	c1.DeploymentStatus = &structs.AllocDeploymentStatus{
  5782  		Healthy: helper.BoolToPtr(true),
  5783  		Canary:  true,
  5784  	}
  5785  
  5786  	// Should still be a canary
  5787  	c2 := mock.Alloc()
  5788  	c2.JobID = j.ID
  5789  	c2.DeploymentID = d.ID
  5790  	d.TaskGroups[c2.TaskGroup].PlacedCanaries = append(d.TaskGroups[c2.TaskGroup].PlacedCanaries, c2.ID)
  5791  	c2.TaskGroup = tg2.Name
  5792  	c2.DeploymentStatus = &structs.AllocDeploymentStatus{
  5793  		Healthy: helper.BoolToPtr(true),
  5794  		Canary:  true,
  5795  	}
  5796  
  5797  	c3 := mock.Alloc()
  5798  	c3.JobID = j.ID
  5799  	c3.DeploymentID = d.ID
  5800  	d.TaskGroups[c3.TaskGroup].PlacedCanaries = append(d.TaskGroups[c3.TaskGroup].PlacedCanaries, c3.ID)
  5801  	c3.DeploymentStatus = &structs.AllocDeploymentStatus{
  5802  		Healthy: helper.BoolToPtr(false),
  5803  		Canary:  true,
  5804  	}
  5805  
  5806  	require.Nil(state.UpsertAllocs(3, []*structs.Allocation{c1, c2, c3}))
  5807  
  5808  	// Create an eval
  5809  	e := mock.Eval()
  5810  
  5811  	// Promote the canaries
  5812  	req := &structs.ApplyDeploymentPromoteRequest{
  5813  		DeploymentPromoteRequest: structs.DeploymentPromoteRequest{
  5814  			DeploymentID: d.ID,
  5815  			Groups:       []string{"web"},
  5816  		},
  5817  		Eval: e,
  5818  	}
  5819  	require.Nil(state.UpdateDeploymentPromotion(4, req))
  5820  
  5821  	// Check that the status per task group was updated properly
  5822  	ws := memdb.NewWatchSet()
  5823  	dout, err := state.DeploymentByID(ws, d.ID)
  5824  	require.Nil(err)
  5825  	require.Len(dout.TaskGroups, 2)
  5826  	require.Contains(dout.TaskGroups, "web")
  5827  	require.True(dout.TaskGroups["web"].Promoted)
  5828  
  5829  	// Check that the evaluation was created
  5830  	eout, err := state.EvalByID(ws, e.ID)
  5831  	require.Nil(err)
  5832  	require.NotNil(eout)
  5833  
  5834  	// Check the canary field was set properly
  5835  	aout1, err1 := state.AllocByID(ws, c1.ID)
  5836  	aout2, err2 := state.AllocByID(ws, c2.ID)
  5837  	aout3, err3 := state.AllocByID(ws, c3.ID)
  5838  	require.Nil(err1)
  5839  	require.Nil(err2)
  5840  	require.Nil(err3)
  5841  	require.NotNil(aout1)
  5842  	require.NotNil(aout2)
  5843  	require.NotNil(aout3)
  5844  	require.False(aout1.DeploymentStatus.Canary)
  5845  	require.True(aout2.DeploymentStatus.Canary)
  5846  	require.True(aout3.DeploymentStatus.Canary)
  5847  }
  5848  
  5849  // Test that allocation health can't be set against a nonexistent deployment
  5850  func TestStateStore_UpsertDeploymentAllocHealth_Nonexistent(t *testing.T) {
  5851  	state := testStateStore(t)
  5852  
  5853  	// Set health against the nonexistent deployment
  5854  	req := &structs.ApplyDeploymentAllocHealthRequest{
  5855  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  5856  			DeploymentID:         uuid.Generate(),
  5857  			HealthyAllocationIDs: []string{uuid.Generate()},
  5858  		},
  5859  	}
  5860  	err := state.UpdateDeploymentAllocHealth(2, req)
  5861  	if err == nil || !strings.Contains(err.Error(), "does not exist") {
  5862  		t.Fatalf("expected error because the deployment doesn't exist: %v", err)
  5863  	}
  5864  }
  5865  
  5866  // Test that allocation health can't be set against a terminal deployment
  5867  func TestStateStore_UpsertDeploymentAllocHealth_Terminal(t *testing.T) {
  5868  	state := testStateStore(t)
  5869  
  5870  	// Insert a terminal deployment
  5871  	d := mock.Deployment()
  5872  	d.Status = structs.DeploymentStatusFailed
  5873  
  5874  	if err := state.UpsertDeployment(1, d); err != nil {
  5875  		t.Fatalf("bad: %v", err)
  5876  	}
  5877  
  5878  	// Set health against the terminal deployment
  5879  	req := &structs.ApplyDeploymentAllocHealthRequest{
  5880  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  5881  			DeploymentID:         d.ID,
  5882  			HealthyAllocationIDs: []string{uuid.Generate()},
  5883  		},
  5884  	}
  5885  	err := state.UpdateDeploymentAllocHealth(2, req)
  5886  	if err == nil || !strings.Contains(err.Error(), "has terminal status") {
  5887  		t.Fatalf("expected error because the deployment is terminal: %v", err)
  5888  	}
  5889  }
  5890  
  5891  // Test that allocation health can't be set against a nonexistent alloc
  5892  func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_Nonexistent(t *testing.T) {
  5893  	state := testStateStore(t)
  5894  
  5895  	// Insert a deployment
  5896  	d := mock.Deployment()
  5897  	if err := state.UpsertDeployment(1, d); err != nil {
  5898  		t.Fatalf("bad: %v", err)
  5899  	}
  5900  
  5901  	// Set health against the terminal deployment
  5902  	req := &structs.ApplyDeploymentAllocHealthRequest{
  5903  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  5904  			DeploymentID:         d.ID,
  5905  			HealthyAllocationIDs: []string{uuid.Generate()},
  5906  		},
  5907  	}
  5908  	err := state.UpdateDeploymentAllocHealth(2, req)
  5909  	if err == nil || !strings.Contains(err.Error(), "unknown alloc") {
  5910  		t.Fatalf("expected error because the alloc doesn't exist: %v", err)
  5911  	}
  5912  }
  5913  
  5914  // Test that allocation health can't be set for an alloc with mismatched
  5915  // deployment ids
  5916  func TestStateStore_UpsertDeploymentAllocHealth_BadAlloc_MismatchDeployment(t *testing.T) {
  5917  	state := testStateStore(t)
  5918  
  5919  	// Insert two  deployment
  5920  	d1 := mock.Deployment()
  5921  	d2 := mock.Deployment()
  5922  	if err := state.UpsertDeployment(1, d1); err != nil {
  5923  		t.Fatalf("bad: %v", err)
  5924  	}
  5925  	if err := state.UpsertDeployment(2, d2); err != nil {
  5926  		t.Fatalf("bad: %v", err)
  5927  	}
  5928  
  5929  	// Insert an alloc for a random deployment
  5930  	a := mock.Alloc()
  5931  	a.DeploymentID = d1.ID
  5932  	if err := state.UpsertAllocs(3, []*structs.Allocation{a}); err != nil {
  5933  		t.Fatalf("bad: %v", err)
  5934  	}
  5935  
  5936  	// Set health against the terminal deployment
  5937  	req := &structs.ApplyDeploymentAllocHealthRequest{
  5938  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  5939  			DeploymentID:         d2.ID,
  5940  			HealthyAllocationIDs: []string{a.ID},
  5941  		},
  5942  	}
  5943  	err := state.UpdateDeploymentAllocHealth(4, req)
  5944  	if err == nil || !strings.Contains(err.Error(), "not part of deployment") {
  5945  		t.Fatalf("expected error because the alloc isn't part of the deployment: %v", err)
  5946  	}
  5947  }
  5948  
  5949  // Test that allocation health is properly set
  5950  func TestStateStore_UpsertDeploymentAllocHealth(t *testing.T) {
  5951  	state := testStateStore(t)
  5952  
  5953  	// Insert a deployment
  5954  	d := mock.Deployment()
  5955  	d.TaskGroups["web"].ProgressDeadline = 5 * time.Minute
  5956  	if err := state.UpsertDeployment(1, d); err != nil {
  5957  		t.Fatalf("bad: %v", err)
  5958  	}
  5959  
  5960  	// Insert two allocations
  5961  	a1 := mock.Alloc()
  5962  	a1.DeploymentID = d.ID
  5963  	a2 := mock.Alloc()
  5964  	a2.DeploymentID = d.ID
  5965  	if err := state.UpsertAllocs(2, []*structs.Allocation{a1, a2}); err != nil {
  5966  		t.Fatalf("bad: %v", err)
  5967  	}
  5968  
  5969  	// Create a job to roll back to
  5970  	j := mock.Job()
  5971  
  5972  	// Create an eval that should be upserted
  5973  	e := mock.Eval()
  5974  
  5975  	// Create a status update for the deployment
  5976  	status, desc := structs.DeploymentStatusFailed, "foo"
  5977  	u := &structs.DeploymentStatusUpdate{
  5978  		DeploymentID:      d.ID,
  5979  		Status:            status,
  5980  		StatusDescription: desc,
  5981  	}
  5982  
  5983  	// Capture the time for the update
  5984  	ts := time.Now()
  5985  
  5986  	// Set health against the deployment
  5987  	req := &structs.ApplyDeploymentAllocHealthRequest{
  5988  		DeploymentAllocHealthRequest: structs.DeploymentAllocHealthRequest{
  5989  			DeploymentID:           d.ID,
  5990  			HealthyAllocationIDs:   []string{a1.ID},
  5991  			UnhealthyAllocationIDs: []string{a2.ID},
  5992  		},
  5993  		Job:              j,
  5994  		Eval:             e,
  5995  		DeploymentUpdate: u,
  5996  		Timestamp:        ts,
  5997  	}
  5998  	err := state.UpdateDeploymentAllocHealth(3, req)
  5999  	if err != nil {
  6000  		t.Fatalf("bad: %v", err)
  6001  	}
  6002  
  6003  	// Check that the status was updated properly
  6004  	ws := memdb.NewWatchSet()
  6005  	dout, err := state.DeploymentByID(ws, d.ID)
  6006  	if err != nil {
  6007  		t.Fatalf("bad: %v", err)
  6008  	}
  6009  	if dout.Status != status || dout.StatusDescription != desc {
  6010  		t.Fatalf("bad: %#v", dout)
  6011  	}
  6012  
  6013  	// Check that the evaluation was created
  6014  	eout, _ := state.EvalByID(ws, e.ID)
  6015  	if err != nil {
  6016  		t.Fatalf("bad: %v", err)
  6017  	}
  6018  	if eout == nil {
  6019  		t.Fatalf("bad: %#v", eout)
  6020  	}
  6021  
  6022  	// Check that the job was created
  6023  	jout, _ := state.JobByID(ws, j.Namespace, j.ID)
  6024  	if err != nil {
  6025  		t.Fatalf("bad: %v", err)
  6026  	}
  6027  	if jout == nil {
  6028  		t.Fatalf("bad: %#v", jout)
  6029  	}
  6030  
  6031  	// Check the status of the allocs
  6032  	out1, err := state.AllocByID(ws, a1.ID)
  6033  	if err != nil {
  6034  		t.Fatalf("err: %v", err)
  6035  	}
  6036  	out2, err := state.AllocByID(ws, a2.ID)
  6037  	if err != nil {
  6038  		t.Fatalf("err: %v", err)
  6039  	}
  6040  
  6041  	if !out1.DeploymentStatus.IsHealthy() {
  6042  		t.Fatalf("bad: alloc %q not healthy", out1.ID)
  6043  	}
  6044  	if !out2.DeploymentStatus.IsUnhealthy() {
  6045  		t.Fatalf("bad: alloc %q not unhealthy", out2.ID)
  6046  	}
  6047  
  6048  	if !out1.DeploymentStatus.Timestamp.Equal(ts) {
  6049  		t.Fatalf("bad: alloc %q had timestamp %v; want %v", out1.ID, out1.DeploymentStatus.Timestamp, ts)
  6050  	}
  6051  	if !out2.DeploymentStatus.Timestamp.Equal(ts) {
  6052  		t.Fatalf("bad: alloc %q had timestamp %v; want %v", out2.ID, out2.DeploymentStatus.Timestamp, ts)
  6053  	}
  6054  }
  6055  
  6056  func TestStateStore_UpsertVaultAccessors(t *testing.T) {
  6057  	state := testStateStore(t)
  6058  	a := mock.VaultAccessor()
  6059  	a2 := mock.VaultAccessor()
  6060  
  6061  	ws := memdb.NewWatchSet()
  6062  	if _, err := state.VaultAccessor(ws, a.Accessor); err != nil {
  6063  		t.Fatalf("err: %v", err)
  6064  	}
  6065  
  6066  	if _, err := state.VaultAccessor(ws, a2.Accessor); err != nil {
  6067  		t.Fatalf("err: %v", err)
  6068  	}
  6069  
  6070  	err := state.UpsertVaultAccessor(1000, []*structs.VaultAccessor{a, a2})
  6071  	if err != nil {
  6072  		t.Fatalf("err: %v", err)
  6073  	}
  6074  
  6075  	if !watchFired(ws) {
  6076  		t.Fatalf("bad")
  6077  	}
  6078  
  6079  	ws = memdb.NewWatchSet()
  6080  	out, err := state.VaultAccessor(ws, a.Accessor)
  6081  	if err != nil {
  6082  		t.Fatalf("err: %v", err)
  6083  	}
  6084  
  6085  	if !reflect.DeepEqual(a, out) {
  6086  		t.Fatalf("bad: %#v %#v", a, out)
  6087  	}
  6088  
  6089  	out, err = state.VaultAccessor(ws, a2.Accessor)
  6090  	if err != nil {
  6091  		t.Fatalf("err: %v", err)
  6092  	}
  6093  
  6094  	if !reflect.DeepEqual(a2, out) {
  6095  		t.Fatalf("bad: %#v %#v", a2, out)
  6096  	}
  6097  
  6098  	iter, err := state.VaultAccessors(ws)
  6099  	if err != nil {
  6100  		t.Fatalf("err: %v", err)
  6101  	}
  6102  
  6103  	count := 0
  6104  	for {
  6105  		raw := iter.Next()
  6106  		if raw == nil {
  6107  			break
  6108  		}
  6109  
  6110  		count++
  6111  		accessor := raw.(*structs.VaultAccessor)
  6112  
  6113  		if !reflect.DeepEqual(accessor, a) && !reflect.DeepEqual(accessor, a2) {
  6114  			t.Fatalf("bad: %#v", accessor)
  6115  		}
  6116  	}
  6117  
  6118  	if count != 2 {
  6119  		t.Fatalf("bad: %d", count)
  6120  	}
  6121  
  6122  	index, err := state.Index("vault_accessors")
  6123  	if err != nil {
  6124  		t.Fatalf("err: %v", err)
  6125  	}
  6126  	if index != 1000 {
  6127  		t.Fatalf("bad: %d", index)
  6128  	}
  6129  
  6130  	if watchFired(ws) {
  6131  		t.Fatalf("bad")
  6132  	}
  6133  }
  6134  
  6135  func TestStateStore_DeleteVaultAccessors(t *testing.T) {
  6136  	state := testStateStore(t)
  6137  	a1 := mock.VaultAccessor()
  6138  	a2 := mock.VaultAccessor()
  6139  	accessors := []*structs.VaultAccessor{a1, a2}
  6140  
  6141  	err := state.UpsertVaultAccessor(1000, accessors)
  6142  	if err != nil {
  6143  		t.Fatalf("err: %v", err)
  6144  	}
  6145  
  6146  	ws := memdb.NewWatchSet()
  6147  	if _, err := state.VaultAccessor(ws, a1.Accessor); err != nil {
  6148  		t.Fatalf("err: %v", err)
  6149  	}
  6150  
  6151  	err = state.DeleteVaultAccessors(1001, accessors)
  6152  	if err != nil {
  6153  		t.Fatalf("err: %v", err)
  6154  	}
  6155  
  6156  	if !watchFired(ws) {
  6157  		t.Fatalf("bad")
  6158  	}
  6159  
  6160  	ws = memdb.NewWatchSet()
  6161  	out, err := state.VaultAccessor(ws, a1.Accessor)
  6162  	if err != nil {
  6163  		t.Fatalf("err: %v", err)
  6164  	}
  6165  	if out != nil {
  6166  		t.Fatalf("bad: %#v %#v", a1, out)
  6167  	}
  6168  	out, err = state.VaultAccessor(ws, a2.Accessor)
  6169  	if err != nil {
  6170  		t.Fatalf("err: %v", err)
  6171  	}
  6172  	if out != nil {
  6173  		t.Fatalf("bad: %#v %#v", a2, out)
  6174  	}
  6175  
  6176  	index, err := state.Index("vault_accessors")
  6177  	if err != nil {
  6178  		t.Fatalf("err: %v", err)
  6179  	}
  6180  	if index != 1001 {
  6181  		t.Fatalf("bad: %d", index)
  6182  	}
  6183  
  6184  	if watchFired(ws) {
  6185  		t.Fatalf("bad")
  6186  	}
  6187  }
  6188  
  6189  func TestStateStore_VaultAccessorsByAlloc(t *testing.T) {
  6190  	state := testStateStore(t)
  6191  	alloc := mock.Alloc()
  6192  	var accessors []*structs.VaultAccessor
  6193  	var expected []*structs.VaultAccessor
  6194  
  6195  	for i := 0; i < 5; i++ {
  6196  		accessor := mock.VaultAccessor()
  6197  		accessor.AllocID = alloc.ID
  6198  		expected = append(expected, accessor)
  6199  		accessors = append(accessors, accessor)
  6200  	}
  6201  
  6202  	for i := 0; i < 10; i++ {
  6203  		accessor := mock.VaultAccessor()
  6204  		accessors = append(accessors, accessor)
  6205  	}
  6206  
  6207  	err := state.UpsertVaultAccessor(1000, accessors)
  6208  	if err != nil {
  6209  		t.Fatalf("err: %v", err)
  6210  	}
  6211  
  6212  	ws := memdb.NewWatchSet()
  6213  	out, err := state.VaultAccessorsByAlloc(ws, alloc.ID)
  6214  	if err != nil {
  6215  		t.Fatalf("err: %v", err)
  6216  	}
  6217  
  6218  	if len(expected) != len(out) {
  6219  		t.Fatalf("bad: %#v %#v", len(expected), len(out))
  6220  	}
  6221  
  6222  	index, err := state.Index("vault_accessors")
  6223  	if err != nil {
  6224  		t.Fatalf("err: %v", err)
  6225  	}
  6226  	if index != 1000 {
  6227  		t.Fatalf("bad: %d", index)
  6228  	}
  6229  
  6230  	if watchFired(ws) {
  6231  		t.Fatalf("bad")
  6232  	}
  6233  }
  6234  
  6235  func TestStateStore_VaultAccessorsByNode(t *testing.T) {
  6236  	state := testStateStore(t)
  6237  	node := mock.Node()
  6238  	var accessors []*structs.VaultAccessor
  6239  	var expected []*structs.VaultAccessor
  6240  
  6241  	for i := 0; i < 5; i++ {
  6242  		accessor := mock.VaultAccessor()
  6243  		accessor.NodeID = node.ID
  6244  		expected = append(expected, accessor)
  6245  		accessors = append(accessors, accessor)
  6246  	}
  6247  
  6248  	for i := 0; i < 10; i++ {
  6249  		accessor := mock.VaultAccessor()
  6250  		accessors = append(accessors, accessor)
  6251  	}
  6252  
  6253  	err := state.UpsertVaultAccessor(1000, accessors)
  6254  	if err != nil {
  6255  		t.Fatalf("err: %v", err)
  6256  	}
  6257  
  6258  	ws := memdb.NewWatchSet()
  6259  	out, err := state.VaultAccessorsByNode(ws, node.ID)
  6260  	if err != nil {
  6261  		t.Fatalf("err: %v", err)
  6262  	}
  6263  
  6264  	if len(expected) != len(out) {
  6265  		t.Fatalf("bad: %#v %#v", len(expected), len(out))
  6266  	}
  6267  
  6268  	index, err := state.Index("vault_accessors")
  6269  	if err != nil {
  6270  		t.Fatalf("err: %v", err)
  6271  	}
  6272  	if index != 1000 {
  6273  		t.Fatalf("bad: %d", index)
  6274  	}
  6275  
  6276  	if watchFired(ws) {
  6277  		t.Fatalf("bad")
  6278  	}
  6279  }
  6280  
  6281  func TestStateStore_RestoreVaultAccessor(t *testing.T) {
  6282  	state := testStateStore(t)
  6283  	a := mock.VaultAccessor()
  6284  
  6285  	restore, err := state.Restore()
  6286  	if err != nil {
  6287  		t.Fatalf("err: %v", err)
  6288  	}
  6289  
  6290  	err = restore.VaultAccessorRestore(a)
  6291  	if err != nil {
  6292  		t.Fatalf("err: %v", err)
  6293  	}
  6294  	restore.Commit()
  6295  
  6296  	ws := memdb.NewWatchSet()
  6297  	out, err := state.VaultAccessor(ws, a.Accessor)
  6298  	if err != nil {
  6299  		t.Fatalf("err: %v", err)
  6300  	}
  6301  
  6302  	if !reflect.DeepEqual(out, a) {
  6303  		t.Fatalf("Bad: %#v %#v", out, a)
  6304  	}
  6305  
  6306  	if watchFired(ws) {
  6307  		t.Fatalf("bad")
  6308  	}
  6309  }
  6310  
  6311  func TestStateStore_UpsertACLPolicy(t *testing.T) {
  6312  	state := testStateStore(t)
  6313  	policy := mock.ACLPolicy()
  6314  	policy2 := mock.ACLPolicy()
  6315  
  6316  	ws := memdb.NewWatchSet()
  6317  	if _, err := state.ACLPolicyByName(ws, policy.Name); err != nil {
  6318  		t.Fatalf("err: %v", err)
  6319  	}
  6320  	if _, err := state.ACLPolicyByName(ws, policy2.Name); err != nil {
  6321  		t.Fatalf("err: %v", err)
  6322  	}
  6323  
  6324  	if err := state.UpsertACLPolicies(1000,
  6325  		[]*structs.ACLPolicy{policy, policy2}); err != nil {
  6326  		t.Fatalf("err: %v", err)
  6327  	}
  6328  	if !watchFired(ws) {
  6329  		t.Fatalf("bad")
  6330  	}
  6331  
  6332  	ws = memdb.NewWatchSet()
  6333  	out, err := state.ACLPolicyByName(ws, policy.Name)
  6334  	assert.Equal(t, nil, err)
  6335  	assert.Equal(t, policy, out)
  6336  
  6337  	out, err = state.ACLPolicyByName(ws, policy2.Name)
  6338  	assert.Equal(t, nil, err)
  6339  	assert.Equal(t, policy2, out)
  6340  
  6341  	iter, err := state.ACLPolicies(ws)
  6342  	if err != nil {
  6343  		t.Fatalf("err: %v", err)
  6344  	}
  6345  
  6346  	// Ensure we see both policies
  6347  	count := 0
  6348  	for {
  6349  		raw := iter.Next()
  6350  		if raw == nil {
  6351  			break
  6352  		}
  6353  		count++
  6354  	}
  6355  	if count != 2 {
  6356  		t.Fatalf("bad: %d", count)
  6357  	}
  6358  
  6359  	index, err := state.Index("acl_policy")
  6360  	if err != nil {
  6361  		t.Fatalf("err: %v", err)
  6362  	}
  6363  	if index != 1000 {
  6364  		t.Fatalf("bad: %d", index)
  6365  	}
  6366  
  6367  	if watchFired(ws) {
  6368  		t.Fatalf("bad")
  6369  	}
  6370  }
  6371  
  6372  func TestStateStore_DeleteACLPolicy(t *testing.T) {
  6373  	state := testStateStore(t)
  6374  	policy := mock.ACLPolicy()
  6375  	policy2 := mock.ACLPolicy()
  6376  
  6377  	// Create the policy
  6378  	if err := state.UpsertACLPolicies(1000,
  6379  		[]*structs.ACLPolicy{policy, policy2}); err != nil {
  6380  		t.Fatalf("err: %v", err)
  6381  	}
  6382  
  6383  	// Create a watcher
  6384  	ws := memdb.NewWatchSet()
  6385  	if _, err := state.ACLPolicyByName(ws, policy.Name); err != nil {
  6386  		t.Fatalf("err: %v", err)
  6387  	}
  6388  
  6389  	// Delete the policy
  6390  	if err := state.DeleteACLPolicies(1001,
  6391  		[]string{policy.Name, policy2.Name}); err != nil {
  6392  		t.Fatalf("err: %v", err)
  6393  	}
  6394  
  6395  	// Ensure watching triggered
  6396  	if !watchFired(ws) {
  6397  		t.Fatalf("bad")
  6398  	}
  6399  
  6400  	// Ensure we don't get the object back
  6401  	ws = memdb.NewWatchSet()
  6402  	out, err := state.ACLPolicyByName(ws, policy.Name)
  6403  	assert.Equal(t, nil, err)
  6404  	if out != nil {
  6405  		t.Fatalf("bad: %#v", out)
  6406  	}
  6407  
  6408  	iter, err := state.ACLPolicies(ws)
  6409  	if err != nil {
  6410  		t.Fatalf("err: %v", err)
  6411  	}
  6412  
  6413  	// Ensure we see both policies
  6414  	count := 0
  6415  	for {
  6416  		raw := iter.Next()
  6417  		if raw == nil {
  6418  			break
  6419  		}
  6420  		count++
  6421  	}
  6422  	if count != 0 {
  6423  		t.Fatalf("bad: %d", count)
  6424  	}
  6425  
  6426  	index, err := state.Index("acl_policy")
  6427  	if err != nil {
  6428  		t.Fatalf("err: %v", err)
  6429  	}
  6430  	if index != 1001 {
  6431  		t.Fatalf("bad: %d", index)
  6432  	}
  6433  
  6434  	if watchFired(ws) {
  6435  		t.Fatalf("bad")
  6436  	}
  6437  }
  6438  
  6439  func TestStateStore_ACLPolicyByNamePrefix(t *testing.T) {
  6440  	state := testStateStore(t)
  6441  	names := []string{
  6442  		"foo",
  6443  		"bar",
  6444  		"foobar",
  6445  		"foozip",
  6446  		"zip",
  6447  	}
  6448  
  6449  	// Create the policies
  6450  	var baseIndex uint64 = 1000
  6451  	for _, name := range names {
  6452  		p := mock.ACLPolicy()
  6453  		p.Name = name
  6454  		if err := state.UpsertACLPolicies(baseIndex, []*structs.ACLPolicy{p}); err != nil {
  6455  			t.Fatalf("err: %v", err)
  6456  		}
  6457  		baseIndex++
  6458  	}
  6459  
  6460  	// Scan by prefix
  6461  	iter, err := state.ACLPolicyByNamePrefix(nil, "foo")
  6462  	if err != nil {
  6463  		t.Fatalf("err: %v", err)
  6464  	}
  6465  
  6466  	// Ensure we see both policies
  6467  	count := 0
  6468  	out := []string{}
  6469  	for {
  6470  		raw := iter.Next()
  6471  		if raw == nil {
  6472  			break
  6473  		}
  6474  		count++
  6475  		out = append(out, raw.(*structs.ACLPolicy).Name)
  6476  	}
  6477  	if count != 3 {
  6478  		t.Fatalf("bad: %d %v", count, out)
  6479  	}
  6480  	sort.Strings(out)
  6481  
  6482  	expect := []string{"foo", "foobar", "foozip"}
  6483  	assert.Equal(t, expect, out)
  6484  }
  6485  
  6486  func TestStateStore_BootstrapACLTokens(t *testing.T) {
  6487  	state := testStateStore(t)
  6488  	tk1 := mock.ACLToken()
  6489  	tk2 := mock.ACLToken()
  6490  
  6491  	ok, resetIdx, err := state.CanBootstrapACLToken()
  6492  	assert.Nil(t, err)
  6493  	assert.Equal(t, true, ok)
  6494  	assert.EqualValues(t, 0, resetIdx)
  6495  
  6496  	if err := state.BootstrapACLTokens(1000, 0, tk1); err != nil {
  6497  		t.Fatalf("err: %v", err)
  6498  	}
  6499  
  6500  	out, err := state.ACLTokenByAccessorID(nil, tk1.AccessorID)
  6501  	assert.Equal(t, nil, err)
  6502  	assert.Equal(t, tk1, out)
  6503  
  6504  	ok, resetIdx, err = state.CanBootstrapACLToken()
  6505  	assert.Nil(t, err)
  6506  	assert.Equal(t, false, ok)
  6507  	assert.EqualValues(t, 1000, resetIdx)
  6508  
  6509  	if err := state.BootstrapACLTokens(1001, 0, tk2); err == nil {
  6510  		t.Fatalf("expected error")
  6511  	}
  6512  
  6513  	iter, err := state.ACLTokens(nil)
  6514  	if err != nil {
  6515  		t.Fatalf("err: %v", err)
  6516  	}
  6517  
  6518  	// Ensure we see both policies
  6519  	count := 0
  6520  	for {
  6521  		raw := iter.Next()
  6522  		if raw == nil {
  6523  			break
  6524  		}
  6525  		count++
  6526  	}
  6527  	if count != 1 {
  6528  		t.Fatalf("bad: %d", count)
  6529  	}
  6530  
  6531  	index, err := state.Index("acl_token")
  6532  	if err != nil {
  6533  		t.Fatalf("err: %v", err)
  6534  	}
  6535  	if index != 1000 {
  6536  		t.Fatalf("bad: %d", index)
  6537  	}
  6538  	index, err = state.Index("acl_token_bootstrap")
  6539  	if err != nil {
  6540  		t.Fatalf("err: %v", err)
  6541  	}
  6542  	if index != 1000 {
  6543  		t.Fatalf("bad: %d", index)
  6544  	}
  6545  
  6546  	// Should allow bootstrap with reset index
  6547  	if err := state.BootstrapACLTokens(1001, 1000, tk2); err != nil {
  6548  		t.Fatalf("err %v", err)
  6549  	}
  6550  
  6551  	// Check we've modified the index
  6552  	index, err = state.Index("acl_token")
  6553  	if err != nil {
  6554  		t.Fatalf("err: %v", err)
  6555  	}
  6556  	if index != 1001 {
  6557  		t.Fatalf("bad: %d", index)
  6558  	}
  6559  	index, err = state.Index("acl_token_bootstrap")
  6560  	if err != nil {
  6561  		t.Fatalf("err: %v", err)
  6562  	}
  6563  	if index != 1001 {
  6564  		t.Fatalf("bad: %d", index)
  6565  	}
  6566  }
  6567  
  6568  func TestStateStore_UpsertACLTokens(t *testing.T) {
  6569  	state := testStateStore(t)
  6570  	tk1 := mock.ACLToken()
  6571  	tk2 := mock.ACLToken()
  6572  
  6573  	ws := memdb.NewWatchSet()
  6574  	if _, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID); err != nil {
  6575  		t.Fatalf("err: %v", err)
  6576  	}
  6577  	if _, err := state.ACLTokenByAccessorID(ws, tk2.AccessorID); err != nil {
  6578  		t.Fatalf("err: %v", err)
  6579  	}
  6580  
  6581  	if err := state.UpsertACLTokens(1000,
  6582  		[]*structs.ACLToken{tk1, tk2}); err != nil {
  6583  		t.Fatalf("err: %v", err)
  6584  	}
  6585  	if !watchFired(ws) {
  6586  		t.Fatalf("bad")
  6587  	}
  6588  
  6589  	ws = memdb.NewWatchSet()
  6590  	out, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID)
  6591  	assert.Equal(t, nil, err)
  6592  	assert.Equal(t, tk1, out)
  6593  
  6594  	out, err = state.ACLTokenByAccessorID(ws, tk2.AccessorID)
  6595  	assert.Equal(t, nil, err)
  6596  	assert.Equal(t, tk2, out)
  6597  
  6598  	out, err = state.ACLTokenBySecretID(ws, tk1.SecretID)
  6599  	assert.Equal(t, nil, err)
  6600  	assert.Equal(t, tk1, out)
  6601  
  6602  	out, err = state.ACLTokenBySecretID(ws, tk2.SecretID)
  6603  	assert.Equal(t, nil, err)
  6604  	assert.Equal(t, tk2, out)
  6605  
  6606  	iter, err := state.ACLTokens(ws)
  6607  	if err != nil {
  6608  		t.Fatalf("err: %v", err)
  6609  	}
  6610  
  6611  	// Ensure we see both policies
  6612  	count := 0
  6613  	for {
  6614  		raw := iter.Next()
  6615  		if raw == nil {
  6616  			break
  6617  		}
  6618  		count++
  6619  	}
  6620  	if count != 2 {
  6621  		t.Fatalf("bad: %d", count)
  6622  	}
  6623  
  6624  	index, err := state.Index("acl_token")
  6625  	if err != nil {
  6626  		t.Fatalf("err: %v", err)
  6627  	}
  6628  	if index != 1000 {
  6629  		t.Fatalf("bad: %d", index)
  6630  	}
  6631  
  6632  	if watchFired(ws) {
  6633  		t.Fatalf("bad")
  6634  	}
  6635  }
  6636  
  6637  func TestStateStore_DeleteACLTokens(t *testing.T) {
  6638  	state := testStateStore(t)
  6639  	tk1 := mock.ACLToken()
  6640  	tk2 := mock.ACLToken()
  6641  
  6642  	// Create the tokens
  6643  	if err := state.UpsertACLTokens(1000,
  6644  		[]*structs.ACLToken{tk1, tk2}); err != nil {
  6645  		t.Fatalf("err: %v", err)
  6646  	}
  6647  
  6648  	// Create a watcher
  6649  	ws := memdb.NewWatchSet()
  6650  	if _, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID); err != nil {
  6651  		t.Fatalf("err: %v", err)
  6652  	}
  6653  
  6654  	// Delete the token
  6655  	if err := state.DeleteACLTokens(1001,
  6656  		[]string{tk1.AccessorID, tk2.AccessorID}); err != nil {
  6657  		t.Fatalf("err: %v", err)
  6658  	}
  6659  
  6660  	// Ensure watching triggered
  6661  	if !watchFired(ws) {
  6662  		t.Fatalf("bad")
  6663  	}
  6664  
  6665  	// Ensure we don't get the object back
  6666  	ws = memdb.NewWatchSet()
  6667  	out, err := state.ACLTokenByAccessorID(ws, tk1.AccessorID)
  6668  	assert.Equal(t, nil, err)
  6669  	if out != nil {
  6670  		t.Fatalf("bad: %#v", out)
  6671  	}
  6672  
  6673  	iter, err := state.ACLTokens(ws)
  6674  	if err != nil {
  6675  		t.Fatalf("err: %v", err)
  6676  	}
  6677  
  6678  	// Ensure we see both policies
  6679  	count := 0
  6680  	for {
  6681  		raw := iter.Next()
  6682  		if raw == nil {
  6683  			break
  6684  		}
  6685  		count++
  6686  	}
  6687  	if count != 0 {
  6688  		t.Fatalf("bad: %d", count)
  6689  	}
  6690  
  6691  	index, err := state.Index("acl_token")
  6692  	if err != nil {
  6693  		t.Fatalf("err: %v", err)
  6694  	}
  6695  	if index != 1001 {
  6696  		t.Fatalf("bad: %d", index)
  6697  	}
  6698  
  6699  	if watchFired(ws) {
  6700  		t.Fatalf("bad")
  6701  	}
  6702  }
  6703  
  6704  func TestStateStore_ACLTokenByAccessorIDPrefix(t *testing.T) {
  6705  	state := testStateStore(t)
  6706  	prefixes := []string{
  6707  		"aaaa",
  6708  		"aabb",
  6709  		"bbbb",
  6710  		"bbcc",
  6711  		"ffff",
  6712  	}
  6713  
  6714  	// Create the tokens
  6715  	var baseIndex uint64 = 1000
  6716  	for _, prefix := range prefixes {
  6717  		tk := mock.ACLToken()
  6718  		tk.AccessorID = prefix + tk.AccessorID[4:]
  6719  		if err := state.UpsertACLTokens(baseIndex, []*structs.ACLToken{tk}); err != nil {
  6720  			t.Fatalf("err: %v", err)
  6721  		}
  6722  		baseIndex++
  6723  	}
  6724  
  6725  	// Scan by prefix
  6726  	iter, err := state.ACLTokenByAccessorIDPrefix(nil, "aa")
  6727  	if err != nil {
  6728  		t.Fatalf("err: %v", err)
  6729  	}
  6730  
  6731  	// Ensure we see both tokens
  6732  	count := 0
  6733  	out := []string{}
  6734  	for {
  6735  		raw := iter.Next()
  6736  		if raw == nil {
  6737  			break
  6738  		}
  6739  		count++
  6740  		out = append(out, raw.(*structs.ACLToken).AccessorID[:4])
  6741  	}
  6742  	if count != 2 {
  6743  		t.Fatalf("bad: %d %v", count, out)
  6744  	}
  6745  	sort.Strings(out)
  6746  
  6747  	expect := []string{"aaaa", "aabb"}
  6748  	assert.Equal(t, expect, out)
  6749  }
  6750  
  6751  func TestStateStore_RestoreACLPolicy(t *testing.T) {
  6752  	state := testStateStore(t)
  6753  	policy := mock.ACLPolicy()
  6754  
  6755  	restore, err := state.Restore()
  6756  	if err != nil {
  6757  		t.Fatalf("err: %v", err)
  6758  	}
  6759  
  6760  	err = restore.ACLPolicyRestore(policy)
  6761  	if err != nil {
  6762  		t.Fatalf("err: %v", err)
  6763  	}
  6764  	restore.Commit()
  6765  
  6766  	ws := memdb.NewWatchSet()
  6767  	out, err := state.ACLPolicyByName(ws, policy.Name)
  6768  	if err != nil {
  6769  		t.Fatalf("err: %v", err)
  6770  	}
  6771  	assert.Equal(t, policy, out)
  6772  }
  6773  
  6774  func TestStateStore_ACLTokensByGlobal(t *testing.T) {
  6775  	state := testStateStore(t)
  6776  	tk1 := mock.ACLToken()
  6777  	tk2 := mock.ACLToken()
  6778  	tk3 := mock.ACLToken()
  6779  	tk4 := mock.ACLToken()
  6780  	tk3.Global = true
  6781  
  6782  	if err := state.UpsertACLTokens(1000,
  6783  		[]*structs.ACLToken{tk1, tk2, tk3, tk4}); err != nil {
  6784  		t.Fatalf("err: %v", err)
  6785  	}
  6786  
  6787  	iter, err := state.ACLTokensByGlobal(nil, true)
  6788  	if err != nil {
  6789  		t.Fatalf("err: %v", err)
  6790  	}
  6791  
  6792  	// Ensure we see the one global policies
  6793  	count := 0
  6794  	for {
  6795  		raw := iter.Next()
  6796  		if raw == nil {
  6797  			break
  6798  		}
  6799  		count++
  6800  	}
  6801  	if count != 1 {
  6802  		t.Fatalf("bad: %d", count)
  6803  	}
  6804  }
  6805  
  6806  func TestStateStore_RestoreACLToken(t *testing.T) {
  6807  	state := testStateStore(t)
  6808  	token := mock.ACLToken()
  6809  
  6810  	restore, err := state.Restore()
  6811  	if err != nil {
  6812  		t.Fatalf("err: %v", err)
  6813  	}
  6814  
  6815  	err = restore.ACLTokenRestore(token)
  6816  	if err != nil {
  6817  		t.Fatalf("err: %v", err)
  6818  	}
  6819  	restore.Commit()
  6820  
  6821  	ws := memdb.NewWatchSet()
  6822  	out, err := state.ACLTokenByAccessorID(ws, token.AccessorID)
  6823  	if err != nil {
  6824  		t.Fatalf("err: %v", err)
  6825  	}
  6826  	assert.Equal(t, token, out)
  6827  }
  6828  
  6829  func TestStateStore_Abandon(t *testing.T) {
  6830  	s := testStateStore(t)
  6831  	abandonCh := s.AbandonCh()
  6832  	s.Abandon()
  6833  	select {
  6834  	case <-abandonCh:
  6835  	default:
  6836  		t.Fatalf("bad")
  6837  	}
  6838  }
  6839  
  6840  // watchFired is a helper for unit tests that returns if the given watch set
  6841  // fired (it doesn't care which watch actually fired). This uses a fixed
  6842  // timeout since we already expect the event happened before calling this and
  6843  // just need to distinguish a fire from a timeout. We do need a little time to
  6844  // allow the watch to set up any goroutines, though.
  6845  func watchFired(ws memdb.WatchSet) bool {
  6846  	timedOut := ws.Watch(time.After(50 * time.Millisecond))
  6847  	return !timedOut
  6848  }
  6849  
  6850  // NodeIDSort is used to sort nodes by ID
  6851  type NodeIDSort []*structs.Node
  6852  
  6853  func (n NodeIDSort) Len() int {
  6854  	return len(n)
  6855  }
  6856  
  6857  func (n NodeIDSort) Less(i, j int) bool {
  6858  	return n[i].ID < n[j].ID
  6859  }
  6860  
  6861  func (n NodeIDSort) Swap(i, j int) {
  6862  	n[i], n[j] = n[j], n[i]
  6863  }
  6864  
  6865  // JobIDis used to sort jobs by id
  6866  type JobIDSort []*structs.Job
  6867  
  6868  func (n JobIDSort) Len() int {
  6869  	return len(n)
  6870  }
  6871  
  6872  func (n JobIDSort) Less(i, j int) bool {
  6873  	return n[i].ID < n[j].ID
  6874  }
  6875  
  6876  func (n JobIDSort) Swap(i, j int) {
  6877  	n[i], n[j] = n[j], n[i]
  6878  }
  6879  
  6880  // EvalIDis used to sort evals by id
  6881  type EvalIDSort []*structs.Evaluation
  6882  
  6883  func (n EvalIDSort) Len() int {
  6884  	return len(n)
  6885  }
  6886  
  6887  func (n EvalIDSort) Less(i, j int) bool {
  6888  	return n[i].ID < n[j].ID
  6889  }
  6890  
  6891  func (n EvalIDSort) Swap(i, j int) {
  6892  	n[i], n[j] = n[j], n[i]
  6893  }
  6894  
  6895  // AllocIDsort used to sort allocations by id
  6896  type AllocIDSort []*structs.Allocation
  6897  
  6898  func (n AllocIDSort) Len() int {
  6899  	return len(n)
  6900  }
  6901  
  6902  func (n AllocIDSort) Less(i, j int) bool {
  6903  	return n[i].ID < n[j].ID
  6904  }
  6905  
  6906  func (n AllocIDSort) Swap(i, j int) {
  6907  	n[i], n[j] = n[j], n[i]
  6908  }