github.com/kaisenlinux/docker.io@v0.0.0-20230510090727-ea55db55fac7/swarmkit/manager/state/store/memory_test.go (about)

     1  package store
     2  
     3  import (
     4  	"errors"
     5  	"strconv"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	events "github.com/docker/go-events"
    11  	"github.com/docker/swarmkit/api"
    12  	"github.com/docker/swarmkit/identity"
    13  	"github.com/docker/swarmkit/manager/state"
    14  	"github.com/docker/swarmkit/manager/state/testutils"
    15  	"github.com/stretchr/testify/assert"
    16  	"github.com/stretchr/testify/require"
    17  )
    18  
    19  var (
    20  	clusterSet = []*api.Cluster{
    21  		{
    22  			ID: "id1",
    23  			Spec: api.ClusterSpec{
    24  				Annotations: api.Annotations{
    25  					Name: "name1",
    26  				},
    27  			},
    28  		},
    29  		{
    30  			ID: "id2",
    31  			Spec: api.ClusterSpec{
    32  				Annotations: api.Annotations{
    33  					Name: "name2",
    34  				},
    35  			},
    36  		},
    37  		{
    38  			ID: "id3",
    39  			Spec: api.ClusterSpec{
    40  				Annotations: api.Annotations{
    41  					Name: "name3",
    42  				},
    43  			},
    44  		},
    45  	}
    46  	altClusterSet = []*api.Cluster{
    47  		{
    48  			ID: "alt-id1",
    49  			Spec: api.ClusterSpec{
    50  				Annotations: api.Annotations{
    51  					Name: "alt-name1",
    52  				},
    53  			},
    54  		},
    55  	}
    56  
    57  	nodeSet = []*api.Node{
    58  		{
    59  			ID: "id1",
    60  			Spec: api.NodeSpec{
    61  				Membership: api.NodeMembershipPending,
    62  			},
    63  			Description: &api.NodeDescription{
    64  				Hostname: "name1",
    65  			},
    66  			Role: api.NodeRoleManager,
    67  		},
    68  		{
    69  			ID: "id2",
    70  			Spec: api.NodeSpec{
    71  				Membership: api.NodeMembershipAccepted,
    72  			},
    73  			Description: &api.NodeDescription{
    74  				Hostname: "name2",
    75  			},
    76  			Role: api.NodeRoleWorker,
    77  		},
    78  		{
    79  			ID: "id3",
    80  			Spec: api.NodeSpec{
    81  				Membership: api.NodeMembershipAccepted,
    82  			},
    83  			Description: &api.NodeDescription{
    84  				// intentionally conflicting hostname
    85  				Hostname: "name2",
    86  			},
    87  			Role: api.NodeRoleWorker,
    88  		},
    89  	}
    90  	altNodeSet = []*api.Node{
    91  		{
    92  			ID: "alt-id1",
    93  			Spec: api.NodeSpec{
    94  				Membership: api.NodeMembershipPending,
    95  			},
    96  			Description: &api.NodeDescription{
    97  				Hostname: "alt-name1",
    98  			},
    99  			Role: api.NodeRoleManager,
   100  		},
   101  	}
   102  
   103  	serviceSet = []*api.Service{
   104  		{
   105  			ID: "id1",
   106  			Spec: api.ServiceSpec{
   107  				Annotations: api.Annotations{
   108  					Name: "name1",
   109  				},
   110  			},
   111  		},
   112  		{
   113  			ID: "id2",
   114  			Spec: api.ServiceSpec{
   115  				Annotations: api.Annotations{
   116  					Name: "name2",
   117  				},
   118  				Mode: &api.ServiceSpec_Global{
   119  					Global: &api.GlobalService{},
   120  				},
   121  			},
   122  		},
   123  		{
   124  			ID: "id3",
   125  			Spec: api.ServiceSpec{
   126  				Annotations: api.Annotations{
   127  					Name: "name3",
   128  				},
   129  			},
   130  		},
   131  	}
   132  	altServiceSet = []*api.Service{
   133  		{
   134  			ID: "alt-id1",
   135  			Spec: api.ServiceSpec{
   136  				Annotations: api.Annotations{
   137  					Name: "alt-name1",
   138  				},
   139  			},
   140  		},
   141  	}
   142  
   143  	taskSet = []*api.Task{
   144  		{
   145  			ID: "id1",
   146  			Annotations: api.Annotations{
   147  				Name: "name1",
   148  			},
   149  			ServiceAnnotations: api.Annotations{
   150  				Name: "name1",
   151  			},
   152  			DesiredState: api.TaskStateRunning,
   153  			NodeID:       nodeSet[0].ID,
   154  		},
   155  		{
   156  			ID: "id2",
   157  			Annotations: api.Annotations{
   158  				Name: "name2.1",
   159  			},
   160  			ServiceAnnotations: api.Annotations{
   161  				Name: "name2",
   162  			},
   163  			DesiredState: api.TaskStateRunning,
   164  			ServiceID:    serviceSet[0].ID,
   165  		},
   166  		{
   167  			ID: "id3",
   168  			Annotations: api.Annotations{
   169  				Name: "name2.2",
   170  			},
   171  			ServiceAnnotations: api.Annotations{
   172  				Name: "name2",
   173  			},
   174  			DesiredState: api.TaskStateShutdown,
   175  		},
   176  	}
   177  	altTaskSet = []*api.Task{
   178  		{
   179  			ID: "alt-id1",
   180  			Annotations: api.Annotations{
   181  				Name: "alt-name1",
   182  			},
   183  			ServiceAnnotations: api.Annotations{
   184  				Name: "alt-name1",
   185  			},
   186  			DesiredState: api.TaskStateRunning,
   187  			NodeID:       altNodeSet[0].ID,
   188  		},
   189  	}
   190  
   191  	networkSet = []*api.Network{
   192  		{
   193  			ID: "id1",
   194  			Spec: api.NetworkSpec{
   195  				Annotations: api.Annotations{
   196  					Name: "name1",
   197  				},
   198  			},
   199  		},
   200  		{
   201  			ID: "id2",
   202  			Spec: api.NetworkSpec{
   203  				Annotations: api.Annotations{
   204  					Name: "name2",
   205  				},
   206  			},
   207  		},
   208  		{
   209  			ID: "id3",
   210  			Spec: api.NetworkSpec{
   211  				Annotations: api.Annotations{
   212  					Name: "name3",
   213  				},
   214  			},
   215  		},
   216  	}
   217  	altNetworkSet = []*api.Network{
   218  		{
   219  			ID: "alt-id1",
   220  			Spec: api.NetworkSpec{
   221  				Annotations: api.Annotations{
   222  					Name: "alt-name1",
   223  				},
   224  			},
   225  		},
   226  	}
   227  
   228  	configSet = []*api.Config{
   229  		{
   230  			ID: "id1",
   231  			Spec: api.ConfigSpec{
   232  				Annotations: api.Annotations{
   233  					Name: "name1",
   234  				},
   235  			},
   236  		},
   237  		{
   238  			ID: "id2",
   239  			Spec: api.ConfigSpec{
   240  				Annotations: api.Annotations{
   241  					Name: "name2",
   242  				},
   243  			},
   244  		},
   245  		{
   246  			ID: "id3",
   247  			Spec: api.ConfigSpec{
   248  				Annotations: api.Annotations{
   249  					Name: "name3",
   250  				},
   251  			},
   252  		},
   253  	}
   254  	altConfigSet = []*api.Config{
   255  		{
   256  			ID: "alt-id1",
   257  			Spec: api.ConfigSpec{
   258  				Annotations: api.Annotations{
   259  					Name: "alt-name1",
   260  				},
   261  			},
   262  		},
   263  	}
   264  
   265  	secretSet = []*api.Secret{
   266  		{
   267  			ID: "id1",
   268  			Spec: api.SecretSpec{
   269  				Annotations: api.Annotations{
   270  					Name: "name1",
   271  				},
   272  			},
   273  		},
   274  		{
   275  			ID: "id2",
   276  			Spec: api.SecretSpec{
   277  				Annotations: api.Annotations{
   278  					Name: "name2",
   279  				},
   280  			},
   281  		},
   282  		{
   283  			ID: "id3",
   284  			Spec: api.SecretSpec{
   285  				Annotations: api.Annotations{
   286  					Name: "name3",
   287  				},
   288  			},
   289  		},
   290  	}
   291  	altSecretSet = []*api.Secret{
   292  		{
   293  			ID: "alt-id1",
   294  			Spec: api.SecretSpec{
   295  				Annotations: api.Annotations{
   296  					Name: "alt-name1",
   297  				},
   298  			},
   299  		},
   300  	}
   301  
   302  	extensionSet = []*api.Extension{
   303  		{
   304  			ID: "id1",
   305  			Annotations: api.Annotations{
   306  				Name: "name1",
   307  			},
   308  		},
   309  		{
   310  			ID: "id2",
   311  			Annotations: api.Annotations{
   312  				Name: "name2",
   313  			},
   314  		},
   315  		{
   316  			ID: "id3",
   317  			Annotations: api.Annotations{
   318  				Name: "name3",
   319  			},
   320  		},
   321  	}
   322  	altExtensionSet = []*api.Extension{
   323  		{
   324  			ID: "alt-id1",
   325  			Annotations: api.Annotations{
   326  				Name: "alt-name1",
   327  			},
   328  		},
   329  	}
   330  
   331  	resourceSet = []*api.Resource{
   332  		{
   333  			ID: "id1",
   334  			Annotations: api.Annotations{
   335  				Name: "name1",
   336  			},
   337  			Kind: "name1", // corresponds to extension id1
   338  		},
   339  		{
   340  			ID: "id2",
   341  			Annotations: api.Annotations{
   342  				Name: "name2",
   343  			},
   344  			Kind: "name2", // corresponds to extension id2
   345  		},
   346  		{
   347  			ID: "id3",
   348  			Annotations: api.Annotations{
   349  				Name: "name3",
   350  			},
   351  			Kind: "name3", // corresponds to extension id3
   352  		},
   353  	}
   354  	altResourceSet = []*api.Resource{
   355  		{
   356  			ID: "alt-id1",
   357  			Annotations: api.Annotations{
   358  				Name: "alt-name1",
   359  			},
   360  			Kind: "alt-name1", // corresponds to extension alt-id1
   361  		},
   362  	}
   363  )
   364  
   365  func setupTestStore(t *testing.T, s *MemoryStore) {
   366  	populateTestStore(t, s,
   367  		clusterSet, nodeSet, serviceSet, taskSet, networkSet, configSet, secretSet,
   368  		extensionSet, resourceSet)
   369  }
   370  
   371  func populateTestStore(t *testing.T, s *MemoryStore,
   372  	clusters []*api.Cluster, nodes []*api.Node, services []*api.Service, tasks []*api.Task, networks []*api.Network,
   373  	configs []*api.Config, secrets []*api.Secret, extensions []*api.Extension, resources []*api.Resource) {
   374  	err := s.Update(func(tx Tx) error {
   375  		// Prepoulate clusters
   376  		for _, c := range clusters {
   377  			assert.NoError(t, CreateCluster(tx, c))
   378  		}
   379  
   380  		// Prepoulate nodes
   381  		for _, n := range nodes {
   382  			assert.NoError(t, CreateNode(tx, n))
   383  		}
   384  
   385  		// Prepopulate services
   386  		for _, s := range services {
   387  			assert.NoError(t, CreateService(tx, s))
   388  		}
   389  		// Prepopulate tasks
   390  		for _, task := range tasks {
   391  			assert.NoError(t, CreateTask(tx, task))
   392  		}
   393  		// Prepopulate networks
   394  		for _, n := range networks {
   395  			assert.NoError(t, CreateNetwork(tx, n))
   396  		}
   397  		// Prepopulate configs
   398  		for _, c := range configs {
   399  			assert.NoError(t, CreateConfig(tx, c))
   400  		}
   401  		// Prepopulate secrets
   402  		for _, s := range secrets {
   403  			assert.NoError(t, CreateSecret(tx, s))
   404  		}
   405  		// Prepopulate extensions
   406  		for _, c := range extensions {
   407  			assert.NoError(t, CreateExtension(tx, c))
   408  		}
   409  		// Prepopulate resources
   410  		for _, s := range resources {
   411  			assert.NoError(t, CreateResource(tx, s))
   412  		}
   413  		return nil
   414  	})
   415  	assert.NoError(t, err)
   416  }
   417  
   418  func TestStoreNode(t *testing.T) {
   419  	s := NewMemoryStore(nil)
   420  	assert.NotNil(t, s)
   421  
   422  	s.View(func(readTx ReadTx) {
   423  		allNodes, err := FindNodes(readTx, All)
   424  		assert.NoError(t, err)
   425  		assert.Empty(t, allNodes)
   426  	})
   427  
   428  	setupTestStore(t, s)
   429  
   430  	err := s.Update(func(tx Tx) error {
   431  		allNodes, err := FindNodes(tx, All)
   432  		assert.NoError(t, err)
   433  		assert.Len(t, allNodes, len(nodeSet))
   434  
   435  		assert.Error(t, CreateNode(tx, nodeSet[0]), "duplicate IDs must be rejected")
   436  		return nil
   437  	})
   438  	assert.NoError(t, err)
   439  
   440  	s.View(func(readTx ReadTx) {
   441  		assert.Equal(t, nodeSet[0], GetNode(readTx, "id1"))
   442  		assert.Equal(t, nodeSet[1], GetNode(readTx, "id2"))
   443  		assert.Equal(t, nodeSet[2], GetNode(readTx, "id3"))
   444  
   445  		foundNodes, err := FindNodes(readTx, ByName("name1"))
   446  		assert.NoError(t, err)
   447  		assert.Len(t, foundNodes, 1)
   448  		foundNodes, err = FindNodes(readTx, ByName("name2"))
   449  		assert.NoError(t, err)
   450  		assert.Len(t, foundNodes, 2)
   451  		foundNodes, err = FindNodes(readTx, Or(ByName("name1"), ByName("name2")))
   452  		assert.NoError(t, err)
   453  		assert.Len(t, foundNodes, 3)
   454  		foundNodes, err = FindNodes(readTx, ByName("invalid"))
   455  		assert.NoError(t, err)
   456  		assert.Len(t, foundNodes, 0)
   457  
   458  		foundNodes, err = FindNodes(readTx, ByIDPrefix("id"))
   459  		assert.NoError(t, err)
   460  		assert.Len(t, foundNodes, 3)
   461  
   462  		foundNodes, err = FindNodes(readTx, ByRole(api.NodeRoleManager))
   463  		assert.NoError(t, err)
   464  		assert.Len(t, foundNodes, 1)
   465  
   466  		foundNodes, err = FindNodes(readTx, ByRole(api.NodeRoleWorker))
   467  		assert.NoError(t, err)
   468  		assert.Len(t, foundNodes, 2)
   469  
   470  		foundNodes, err = FindNodes(readTx, ByMembership(api.NodeMembershipPending))
   471  		assert.NoError(t, err)
   472  		assert.Len(t, foundNodes, 1)
   473  
   474  		foundNodes, err = FindNodes(readTx, ByMembership(api.NodeMembershipAccepted))
   475  		assert.NoError(t, err)
   476  		assert.Len(t, foundNodes, 2)
   477  	})
   478  
   479  	// Update.
   480  	update := &api.Node{
   481  		ID: "id3",
   482  		Description: &api.NodeDescription{
   483  			Hostname: "name3",
   484  		},
   485  	}
   486  	err = s.Update(func(tx Tx) error {
   487  		assert.NotEqual(t, update, GetNode(tx, "id3"))
   488  		assert.NoError(t, UpdateNode(tx, update))
   489  		assert.Equal(t, update, GetNode(tx, "id3"))
   490  
   491  		foundNodes, err := FindNodes(tx, ByName("name2"))
   492  		assert.NoError(t, err)
   493  		assert.Len(t, foundNodes, 1)
   494  		foundNodes, err = FindNodes(tx, ByName("name3"))
   495  		assert.NoError(t, err)
   496  		assert.Len(t, foundNodes, 1)
   497  
   498  		invalidUpdate := *nodeSet[0]
   499  		invalidUpdate.ID = "invalid"
   500  		assert.Error(t, UpdateNode(tx, &invalidUpdate), "invalid IDs should be rejected")
   501  
   502  		// Delete
   503  		assert.NotNil(t, GetNode(tx, "id1"))
   504  		assert.NoError(t, DeleteNode(tx, "id1"))
   505  		assert.Nil(t, GetNode(tx, "id1"))
   506  		foundNodes, err = FindNodes(tx, ByName("name1"))
   507  		assert.NoError(t, err)
   508  		assert.Empty(t, foundNodes)
   509  
   510  		assert.Equal(t, DeleteNode(tx, "nonexistent"), ErrNotExist)
   511  		return nil
   512  	})
   513  	assert.NoError(t, err)
   514  }
   515  
   516  func TestStoreService(t *testing.T) {
   517  	s := NewMemoryStore(nil)
   518  	assert.NotNil(t, s)
   519  
   520  	s.View(func(readTx ReadTx) {
   521  		allServices, err := FindServices(readTx, All)
   522  		assert.NoError(t, err)
   523  		assert.Empty(t, allServices)
   524  	})
   525  
   526  	setupTestStore(t, s)
   527  
   528  	err := s.Update(func(tx Tx) error {
   529  		assert.Equal(t,
   530  			CreateService(tx, &api.Service{
   531  				ID: "id1",
   532  				Spec: api.ServiceSpec{
   533  					Annotations: api.Annotations{
   534  						Name: "name4",
   535  					},
   536  				},
   537  			}), ErrExist, "duplicate IDs must be rejected")
   538  
   539  		assert.Equal(t,
   540  			CreateService(tx, &api.Service{
   541  				ID: "id4",
   542  				Spec: api.ServiceSpec{
   543  					Annotations: api.Annotations{
   544  						Name: "name1",
   545  					},
   546  				},
   547  			}), ErrNameConflict, "duplicate names must be rejected")
   548  
   549  		assert.Equal(t,
   550  			CreateService(tx, &api.Service{
   551  				ID: "id4",
   552  				Spec: api.ServiceSpec{
   553  					Annotations: api.Annotations{
   554  						Name: "NAME1",
   555  					},
   556  				},
   557  			}), ErrNameConflict, "duplicate check should be case insensitive")
   558  		return nil
   559  	})
   560  	assert.NoError(t, err)
   561  
   562  	s.View(func(readTx ReadTx) {
   563  		assert.Equal(t, serviceSet[0], GetService(readTx, "id1"))
   564  		assert.Equal(t, serviceSet[1], GetService(readTx, "id2"))
   565  		assert.Equal(t, serviceSet[2], GetService(readTx, "id3"))
   566  
   567  		foundServices, err := FindServices(readTx, ByNamePrefix("name1"))
   568  		assert.NoError(t, err)
   569  		assert.Len(t, foundServices, 1)
   570  		foundServices, err = FindServices(readTx, ByNamePrefix("NAME1"))
   571  		assert.NoError(t, err)
   572  		assert.Len(t, foundServices, 1)
   573  		foundServices, err = FindServices(readTx, ByNamePrefix("invalid"))
   574  		assert.NoError(t, err)
   575  		assert.Len(t, foundServices, 0)
   576  		foundServices, err = FindServices(readTx, Or(ByNamePrefix("name1"), ByNamePrefix("name2")))
   577  		assert.NoError(t, err)
   578  		assert.Len(t, foundServices, 2)
   579  		foundServices, err = FindServices(readTx, Or(ByNamePrefix("name1"), ByNamePrefix("name2"), ByNamePrefix("name4")))
   580  		assert.NoError(t, err)
   581  		assert.Len(t, foundServices, 2)
   582  
   583  		foundServices, err = FindServices(readTx, ByIDPrefix("id"))
   584  		assert.NoError(t, err)
   585  		assert.Len(t, foundServices, 3)
   586  	})
   587  
   588  	// Update.
   589  	err = s.Update(func(tx Tx) error {
   590  		// Regular update.
   591  		update := serviceSet[0].Copy()
   592  		update.Spec.Annotations.Labels = map[string]string{
   593  			"foo": "bar",
   594  		}
   595  
   596  		assert.NotEqual(t, update, GetService(tx, update.ID))
   597  		assert.NoError(t, UpdateService(tx, update))
   598  		assert.Equal(t, update, GetService(tx, update.ID))
   599  
   600  		// Name conflict.
   601  		update = GetService(tx, update.ID)
   602  		update.Spec.Annotations.Name = "name2"
   603  		assert.Equal(t, UpdateService(tx, update), ErrNameConflict, "duplicate names should be rejected")
   604  		update = GetService(tx, update.ID)
   605  		update.Spec.Annotations.Name = "NAME2"
   606  		assert.Equal(t, UpdateService(tx, update), ErrNameConflict, "duplicate check should be case insensitive")
   607  
   608  		// Name change.
   609  		update = GetService(tx, update.ID)
   610  		foundServices, err := FindServices(tx, ByNamePrefix("name1"))
   611  		assert.NoError(t, err)
   612  		assert.Len(t, foundServices, 1)
   613  		foundServices, err = FindServices(tx, ByNamePrefix("name4"))
   614  		assert.NoError(t, err)
   615  		assert.Empty(t, foundServices)
   616  
   617  		update.Spec.Annotations.Name = "name4"
   618  		assert.NoError(t, UpdateService(tx, update))
   619  		foundServices, err = FindServices(tx, ByNamePrefix("name1"))
   620  		assert.NoError(t, err)
   621  		assert.Empty(t, foundServices)
   622  		foundServices, err = FindServices(tx, ByNamePrefix("name4"))
   623  		assert.NoError(t, err)
   624  		assert.Len(t, foundServices, 1)
   625  
   626  		// Invalid update.
   627  		invalidUpdate := serviceSet[0].Copy()
   628  		invalidUpdate.ID = "invalid"
   629  		assert.Error(t, UpdateService(tx, invalidUpdate), "invalid IDs should be rejected")
   630  
   631  		return nil
   632  	})
   633  	assert.NoError(t, err)
   634  
   635  	// Delete
   636  	err = s.Update(func(tx Tx) error {
   637  		assert.NotNil(t, GetService(tx, "id1"))
   638  		assert.NoError(t, DeleteService(tx, "id1"))
   639  		assert.Nil(t, GetService(tx, "id1"))
   640  		foundServices, err := FindServices(tx, ByNamePrefix("name1"))
   641  		assert.NoError(t, err)
   642  		assert.Empty(t, foundServices)
   643  
   644  		assert.Equal(t, DeleteService(tx, "nonexistent"), ErrNotExist)
   645  		return nil
   646  	})
   647  	assert.NoError(t, err)
   648  }
   649  
   650  func TestStoreNetwork(t *testing.T) {
   651  	s := NewMemoryStore(nil)
   652  	assert.NotNil(t, s)
   653  
   654  	s.View(func(readTx ReadTx) {
   655  		allNetworks, err := FindNetworks(readTx, All)
   656  		assert.NoError(t, err)
   657  		assert.Empty(t, allNetworks)
   658  	})
   659  
   660  	setupTestStore(t, s)
   661  
   662  	err := s.Update(func(tx Tx) error {
   663  		allNetworks, err := FindNetworks(tx, All)
   664  		assert.NoError(t, err)
   665  		assert.Len(t, allNetworks, len(networkSet))
   666  
   667  		assert.Error(t, CreateNetwork(tx, networkSet[0]), "duplicate IDs must be rejected")
   668  		return nil
   669  	})
   670  	assert.NoError(t, err)
   671  
   672  	s.View(func(readTx ReadTx) {
   673  		assert.Equal(t, networkSet[0], GetNetwork(readTx, "id1"))
   674  		assert.Equal(t, networkSet[1], GetNetwork(readTx, "id2"))
   675  		assert.Equal(t, networkSet[2], GetNetwork(readTx, "id3"))
   676  
   677  		foundNetworks, err := FindNetworks(readTx, ByName("name1"))
   678  		assert.NoError(t, err)
   679  		assert.Len(t, foundNetworks, 1)
   680  		foundNetworks, err = FindNetworks(readTx, ByName("name2"))
   681  		assert.NoError(t, err)
   682  		assert.Len(t, foundNetworks, 1)
   683  		foundNetworks, err = FindNetworks(readTx, ByName("invalid"))
   684  		assert.NoError(t, err)
   685  		assert.Len(t, foundNetworks, 0)
   686  	})
   687  
   688  	err = s.Update(func(tx Tx) error {
   689  		// Delete
   690  		assert.NotNil(t, GetNetwork(tx, "id1"))
   691  		assert.NoError(t, DeleteNetwork(tx, "id1"))
   692  		assert.Nil(t, GetNetwork(tx, "id1"))
   693  		foundNetworks, err := FindNetworks(tx, ByName("name1"))
   694  		assert.NoError(t, err)
   695  		assert.Empty(t, foundNetworks)
   696  
   697  		assert.Equal(t, DeleteNetwork(tx, "nonexistent"), ErrNotExist)
   698  		return nil
   699  	})
   700  
   701  	assert.NoError(t, err)
   702  }
   703  
   704  func TestStoreTask(t *testing.T) {
   705  	s := NewMemoryStore(nil)
   706  	assert.NotNil(t, s)
   707  
   708  	s.View(func(tx ReadTx) {
   709  		allTasks, err := FindTasks(tx, All)
   710  		assert.NoError(t, err)
   711  		assert.Empty(t, allTasks)
   712  	})
   713  
   714  	setupTestStore(t, s)
   715  
   716  	err := s.Update(func(tx Tx) error {
   717  		allTasks, err := FindTasks(tx, All)
   718  		assert.NoError(t, err)
   719  		assert.Len(t, allTasks, len(taskSet))
   720  
   721  		assert.Error(t, CreateTask(tx, taskSet[0]), "duplicate IDs must be rejected")
   722  		return nil
   723  	})
   724  	assert.NoError(t, err)
   725  
   726  	s.View(func(readTx ReadTx) {
   727  		assert.Equal(t, taskSet[0], GetTask(readTx, "id1"))
   728  		assert.Equal(t, taskSet[1], GetTask(readTx, "id2"))
   729  		assert.Equal(t, taskSet[2], GetTask(readTx, "id3"))
   730  
   731  		foundTasks, err := FindTasks(readTx, ByNamePrefix("name1"))
   732  		assert.NoError(t, err)
   733  		assert.Len(t, foundTasks, 1)
   734  		foundTasks, err = FindTasks(readTx, ByNamePrefix("name2"))
   735  		assert.NoError(t, err)
   736  		assert.Len(t, foundTasks, 2)
   737  		foundTasks, err = FindTasks(readTx, ByNamePrefix("invalid"))
   738  		assert.NoError(t, err)
   739  		assert.Len(t, foundTasks, 0)
   740  
   741  		foundTasks, err = FindTasks(readTx, ByNodeID(nodeSet[0].ID))
   742  		assert.NoError(t, err)
   743  		assert.Len(t, foundTasks, 1)
   744  		assert.Equal(t, foundTasks[0], taskSet[0])
   745  		foundTasks, err = FindTasks(readTx, ByNodeID("invalid"))
   746  		assert.NoError(t, err)
   747  		assert.Len(t, foundTasks, 0)
   748  
   749  		foundTasks, err = FindTasks(readTx, ByServiceID(serviceSet[0].ID))
   750  		assert.NoError(t, err)
   751  		assert.Len(t, foundTasks, 1)
   752  		assert.Equal(t, foundTasks[0], taskSet[1])
   753  		foundTasks, err = FindTasks(readTx, ByServiceID("invalid"))
   754  		assert.NoError(t, err)
   755  		assert.Len(t, foundTasks, 0)
   756  
   757  		foundTasks, err = FindTasks(readTx, ByDesiredState(api.TaskStateRunning))
   758  		assert.NoError(t, err)
   759  		assert.Len(t, foundTasks, 2)
   760  		assert.Equal(t, foundTasks[0].DesiredState, api.TaskStateRunning)
   761  		assert.Equal(t, foundTasks[0].DesiredState, api.TaskStateRunning)
   762  		foundTasks, err = FindTasks(readTx, ByDesiredState(api.TaskStateShutdown))
   763  		assert.NoError(t, err)
   764  		assert.Len(t, foundTasks, 1)
   765  		assert.Equal(t, foundTasks[0], taskSet[2])
   766  		foundTasks, err = FindTasks(readTx, ByDesiredState(api.TaskStatePending))
   767  		assert.NoError(t, err)
   768  		assert.Len(t, foundTasks, 0)
   769  	})
   770  
   771  	// Update.
   772  	update := &api.Task{
   773  		ID: "id3",
   774  		Annotations: api.Annotations{
   775  			Name: "name3",
   776  		},
   777  		ServiceAnnotations: api.Annotations{
   778  			Name: "name3",
   779  		},
   780  	}
   781  	err = s.Update(func(tx Tx) error {
   782  		assert.NotEqual(t, update, GetTask(tx, "id3"))
   783  		assert.NoError(t, UpdateTask(tx, update))
   784  		assert.Equal(t, update, GetTask(tx, "id3"))
   785  
   786  		foundTasks, err := FindTasks(tx, ByNamePrefix("name2"))
   787  		assert.NoError(t, err)
   788  		assert.Len(t, foundTasks, 1)
   789  		foundTasks, err = FindTasks(tx, ByNamePrefix("name3"))
   790  		assert.NoError(t, err)
   791  		assert.Len(t, foundTasks, 1)
   792  
   793  		invalidUpdate := *taskSet[0]
   794  		invalidUpdate.ID = "invalid"
   795  		assert.Error(t, UpdateTask(tx, &invalidUpdate), "invalid IDs should be rejected")
   796  
   797  		// Delete
   798  		assert.NotNil(t, GetTask(tx, "id1"))
   799  		assert.NoError(t, DeleteTask(tx, "id1"))
   800  		assert.Nil(t, GetTask(tx, "id1"))
   801  		foundTasks, err = FindTasks(tx, ByNamePrefix("name1"))
   802  		assert.NoError(t, err)
   803  		assert.Empty(t, foundTasks)
   804  
   805  		assert.Equal(t, DeleteTask(tx, "nonexistent"), ErrNotExist)
   806  		return nil
   807  	})
   808  	assert.NoError(t, err)
   809  }
   810  
   811  func TestStoreSnapshot(t *testing.T) {
   812  	s1 := NewMemoryStore(nil)
   813  	assert.NotNil(t, s1)
   814  
   815  	setupTestStore(t, s1)
   816  
   817  	s2 := NewMemoryStore(nil)
   818  	assert.NotNil(t, s2)
   819  
   820  	copyToS2 := func(readTx ReadTx) error {
   821  		return s2.Update(func(tx Tx) error {
   822  			// Copy over new data
   823  			nodes, err := FindNodes(readTx, All)
   824  			if err != nil {
   825  				return err
   826  			}
   827  			for _, n := range nodes {
   828  				if err := CreateNode(tx, n); err != nil {
   829  					return err
   830  				}
   831  			}
   832  
   833  			tasks, err := FindTasks(readTx, All)
   834  			if err != nil {
   835  				return err
   836  			}
   837  			for _, t := range tasks {
   838  				if err := CreateTask(tx, t); err != nil {
   839  					return err
   840  				}
   841  			}
   842  
   843  			services, err := FindServices(readTx, All)
   844  			if err != nil {
   845  				return err
   846  			}
   847  			for _, s := range services {
   848  				if err := CreateService(tx, s); err != nil {
   849  					return err
   850  				}
   851  			}
   852  
   853  			networks, err := FindNetworks(readTx, All)
   854  			if err != nil {
   855  				return err
   856  			}
   857  			for _, n := range networks {
   858  				if err := CreateNetwork(tx, n); err != nil {
   859  					return err
   860  				}
   861  			}
   862  
   863  			return nil
   864  		})
   865  	}
   866  
   867  	// Fork
   868  	watcher, cancel, err := ViewAndWatch(s1, copyToS2)
   869  	defer cancel()
   870  	assert.NoError(t, err)
   871  
   872  	s2.View(func(tx2 ReadTx) {
   873  		assert.Equal(t, nodeSet[0], GetNode(tx2, "id1"))
   874  		assert.Equal(t, nodeSet[1], GetNode(tx2, "id2"))
   875  		assert.Equal(t, nodeSet[2], GetNode(tx2, "id3"))
   876  
   877  		assert.Equal(t, serviceSet[0], GetService(tx2, "id1"))
   878  		assert.Equal(t, serviceSet[1], GetService(tx2, "id2"))
   879  		assert.Equal(t, serviceSet[2], GetService(tx2, "id3"))
   880  
   881  		assert.Equal(t, taskSet[0], GetTask(tx2, "id1"))
   882  		assert.Equal(t, taskSet[1], GetTask(tx2, "id2"))
   883  		assert.Equal(t, taskSet[2], GetTask(tx2, "id3"))
   884  	})
   885  
   886  	// Create node
   887  	createNode := &api.Node{
   888  		ID: "id4",
   889  		Spec: api.NodeSpec{
   890  			Annotations: api.Annotations{
   891  				Name: "name4",
   892  			},
   893  		},
   894  	}
   895  
   896  	err = s1.Update(func(tx1 Tx) error {
   897  		assert.NoError(t, CreateNode(tx1, createNode))
   898  		return nil
   899  	})
   900  	assert.NoError(t, err)
   901  
   902  	assert.NoError(t, Apply(s2, <-watcher))
   903  	<-watcher // consume commit event
   904  
   905  	s2.View(func(tx2 ReadTx) {
   906  		assert.Equal(t, createNode, GetNode(tx2, "id4"))
   907  	})
   908  
   909  	// Update node
   910  	updateNode := &api.Node{
   911  		ID: "id3",
   912  		Spec: api.NodeSpec{
   913  			Annotations: api.Annotations{
   914  				Name: "name3",
   915  			},
   916  		},
   917  	}
   918  
   919  	err = s1.Update(func(tx1 Tx) error {
   920  		assert.NoError(t, UpdateNode(tx1, updateNode))
   921  		return nil
   922  	})
   923  	assert.NoError(t, err)
   924  
   925  	assert.NoError(t, Apply(s2, <-watcher))
   926  	<-watcher // consume commit event
   927  
   928  	s2.View(func(tx2 ReadTx) {
   929  		assert.Equal(t, updateNode, GetNode(tx2, "id3"))
   930  	})
   931  
   932  	err = s1.Update(func(tx1 Tx) error {
   933  		// Delete node
   934  		assert.NoError(t, DeleteNode(tx1, "id1"))
   935  		return nil
   936  	})
   937  	assert.NoError(t, err)
   938  
   939  	assert.NoError(t, Apply(s2, <-watcher))
   940  	<-watcher // consume commit event
   941  
   942  	s2.View(func(tx2 ReadTx) {
   943  		assert.Nil(t, GetNode(tx2, "id1"))
   944  	})
   945  
   946  	// Create service
   947  	createService := &api.Service{
   948  		ID: "id4",
   949  		Spec: api.ServiceSpec{
   950  			Annotations: api.Annotations{
   951  				Name: "name4",
   952  			},
   953  		},
   954  	}
   955  
   956  	err = s1.Update(func(tx1 Tx) error {
   957  		assert.NoError(t, CreateService(tx1, createService))
   958  		return nil
   959  	})
   960  	assert.NoError(t, err)
   961  
   962  	assert.NoError(t, Apply(s2, <-watcher))
   963  	<-watcher // consume commit event
   964  
   965  	s2.View(func(tx2 ReadTx) {
   966  		assert.Equal(t, createService, GetService(tx2, "id4"))
   967  	})
   968  
   969  	// Update service
   970  	updateService := serviceSet[2].Copy()
   971  	updateService.Spec.Annotations.Name = "new-name"
   972  	err = s1.Update(func(tx1 Tx) error {
   973  		assert.NotEqual(t, updateService, GetService(tx1, updateService.ID))
   974  		assert.NoError(t, UpdateService(tx1, updateService))
   975  		return nil
   976  	})
   977  	assert.NoError(t, err)
   978  
   979  	assert.NoError(t, Apply(s2, <-watcher))
   980  	<-watcher // consume commit event
   981  
   982  	s2.View(func(tx2 ReadTx) {
   983  		assert.Equal(t, updateService, GetService(tx2, "id3"))
   984  	})
   985  
   986  	err = s1.Update(func(tx1 Tx) error {
   987  		// Delete service
   988  		assert.NoError(t, DeleteService(tx1, "id1"))
   989  		return nil
   990  	})
   991  	assert.NoError(t, err)
   992  
   993  	assert.NoError(t, Apply(s2, <-watcher))
   994  	<-watcher // consume commit event
   995  
   996  	s2.View(func(tx2 ReadTx) {
   997  		assert.Nil(t, GetService(tx2, "id1"))
   998  	})
   999  
  1000  	// Create task
  1001  	createTask := &api.Task{
  1002  		ID: "id4",
  1003  		ServiceAnnotations: api.Annotations{
  1004  			Name: "name4",
  1005  		},
  1006  	}
  1007  
  1008  	err = s1.Update(func(tx1 Tx) error {
  1009  		assert.NoError(t, CreateTask(tx1, createTask))
  1010  		return nil
  1011  	})
  1012  	assert.NoError(t, err)
  1013  
  1014  	assert.NoError(t, Apply(s2, <-watcher))
  1015  	<-watcher // consume commit event
  1016  
  1017  	s2.View(func(tx2 ReadTx) {
  1018  		assert.Equal(t, createTask, GetTask(tx2, "id4"))
  1019  	})
  1020  
  1021  	// Update task
  1022  	updateTask := &api.Task{
  1023  		ID: "id3",
  1024  		ServiceAnnotations: api.Annotations{
  1025  			Name: "name3",
  1026  		},
  1027  	}
  1028  
  1029  	err = s1.Update(func(tx1 Tx) error {
  1030  		assert.NoError(t, UpdateTask(tx1, updateTask))
  1031  		return nil
  1032  	})
  1033  	assert.NoError(t, err)
  1034  	assert.NoError(t, Apply(s2, <-watcher))
  1035  	<-watcher // consume commit event
  1036  
  1037  	s2.View(func(tx2 ReadTx) {
  1038  		assert.Equal(t, updateTask, GetTask(tx2, "id3"))
  1039  	})
  1040  
  1041  	err = s1.Update(func(tx1 Tx) error {
  1042  		// Delete task
  1043  		assert.NoError(t, DeleteTask(tx1, "id1"))
  1044  		return nil
  1045  	})
  1046  	assert.NoError(t, err)
  1047  	assert.NoError(t, Apply(s2, <-watcher))
  1048  	<-watcher // consume commit event
  1049  
  1050  	s2.View(func(tx2 ReadTx) {
  1051  		assert.Nil(t, GetTask(tx2, "id1"))
  1052  	})
  1053  }
  1054  
  1055  func TestCustomIndex(t *testing.T) {
  1056  	s := NewMemoryStore(nil)
  1057  	assert.NotNil(t, s)
  1058  
  1059  	setupTestStore(t, s)
  1060  
  1061  	// Add a custom index entry to each node
  1062  	err := s.Update(func(tx Tx) error {
  1063  		allNodes, err := FindNodes(tx, All)
  1064  		assert.NoError(t, err)
  1065  		assert.Len(t, allNodes, len(nodeSet))
  1066  
  1067  		for _, n := range allNodes {
  1068  			switch n.ID {
  1069  			case "id2":
  1070  				n.Spec.Annotations.Indices = []api.IndexEntry{
  1071  					{Key: "nodesbefore", Val: "id1"},
  1072  				}
  1073  				assert.NoError(t, UpdateNode(tx, n))
  1074  			case "id3":
  1075  				n.Spec.Annotations.Indices = []api.IndexEntry{
  1076  					{Key: "nodesbefore", Val: "id1"},
  1077  					{Key: "nodesbefore", Val: "id2"},
  1078  				}
  1079  				assert.NoError(t, UpdateNode(tx, n))
  1080  			}
  1081  		}
  1082  		return nil
  1083  	})
  1084  	assert.NoError(t, err)
  1085  
  1086  	s.View(func(readTx ReadTx) {
  1087  		foundNodes, err := FindNodes(readTx, ByCustom("", "nodesbefore", "id2"))
  1088  		require.NoError(t, err)
  1089  		require.Len(t, foundNodes, 1)
  1090  		assert.Equal(t, "id3", foundNodes[0].ID)
  1091  
  1092  		foundNodes, err = FindNodes(readTx, ByCustom("", "nodesbefore", "id1"))
  1093  		require.NoError(t, err)
  1094  		require.Len(t, foundNodes, 2)
  1095  
  1096  		foundNodes, err = FindNodes(readTx, ByCustom("", "nodesbefore", "id3"))
  1097  		require.NoError(t, err)
  1098  		require.Len(t, foundNodes, 0)
  1099  
  1100  		foundNodes, err = FindNodes(readTx, ByCustomPrefix("", "nodesbefore", "id"))
  1101  		require.NoError(t, err)
  1102  		require.Len(t, foundNodes, 2)
  1103  
  1104  		foundNodes, err = FindNodes(readTx, ByCustomPrefix("", "nodesbefore", "id6"))
  1105  		require.NoError(t, err)
  1106  		require.Len(t, foundNodes, 0)
  1107  	})
  1108  }
  1109  
  1110  func TestFailedTransaction(t *testing.T) {
  1111  	s := NewMemoryStore(nil)
  1112  	assert.NotNil(t, s)
  1113  
  1114  	// Create one node
  1115  	err := s.Update(func(tx Tx) error {
  1116  		n := &api.Node{
  1117  			ID: "id1",
  1118  			Description: &api.NodeDescription{
  1119  				Hostname: "name1",
  1120  			},
  1121  		}
  1122  
  1123  		assert.NoError(t, CreateNode(tx, n))
  1124  		return nil
  1125  	})
  1126  	assert.NoError(t, err)
  1127  
  1128  	// Create a second node, but then roll back the transaction
  1129  	err = s.Update(func(tx Tx) error {
  1130  		n := &api.Node{
  1131  			ID: "id2",
  1132  			Description: &api.NodeDescription{
  1133  				Hostname: "name2",
  1134  			},
  1135  		}
  1136  
  1137  		assert.NoError(t, CreateNode(tx, n))
  1138  		return errors.New("rollback")
  1139  	})
  1140  	assert.Error(t, err)
  1141  
  1142  	s.View(func(tx ReadTx) {
  1143  		foundNodes, err := FindNodes(tx, All)
  1144  		assert.NoError(t, err)
  1145  		assert.Len(t, foundNodes, 1)
  1146  		foundNodes, err = FindNodes(tx, ByName("name1"))
  1147  		assert.NoError(t, err)
  1148  		assert.Len(t, foundNodes, 1)
  1149  		foundNodes, err = FindNodes(tx, ByName("name2"))
  1150  		assert.NoError(t, err)
  1151  		assert.Len(t, foundNodes, 0)
  1152  	})
  1153  }
  1154  
  1155  func TestVersion(t *testing.T) {
  1156  	s := NewMemoryStore(&testutils.MockProposer{})
  1157  	assert.NotNil(t, s)
  1158  
  1159  	var (
  1160  		retrievedNode  *api.Node
  1161  		retrievedNode2 *api.Node
  1162  	)
  1163  
  1164  	// Create one node
  1165  	n := &api.Node{
  1166  		ID: "id1",
  1167  		Spec: api.NodeSpec{
  1168  			Annotations: api.Annotations{
  1169  				Name: "name1",
  1170  			},
  1171  		},
  1172  	}
  1173  	err := s.Update(func(tx Tx) error {
  1174  		assert.NoError(t, CreateNode(tx, n))
  1175  		return nil
  1176  	})
  1177  	assert.NoError(t, err)
  1178  
  1179  	// Update the node using an object fetched from the store.
  1180  	n.Spec.Annotations.Name = "name2"
  1181  	err = s.Update(func(tx Tx) error {
  1182  		assert.NoError(t, UpdateNode(tx, n))
  1183  		retrievedNode = GetNode(tx, n.ID)
  1184  		return nil
  1185  	})
  1186  	assert.NoError(t, err)
  1187  
  1188  	// Make sure the store is updating our local copy with the version.
  1189  	assert.Equal(t, n.Meta.Version, retrievedNode.Meta.Version)
  1190  
  1191  	// Try again, this time using the retrieved node.
  1192  	retrievedNode.Spec.Annotations.Name = "name2"
  1193  	err = s.Update(func(tx Tx) error {
  1194  		assert.NoError(t, UpdateNode(tx, retrievedNode))
  1195  		retrievedNode2 = GetNode(tx, n.ID)
  1196  		return nil
  1197  	})
  1198  	assert.NoError(t, err)
  1199  
  1200  	// Try to update retrievedNode again. This should fail because it was
  1201  	// already used to perform an update.
  1202  	retrievedNode.Spec.Annotations.Name = "name3"
  1203  	err = s.Update(func(tx Tx) error {
  1204  		assert.Equal(t, ErrSequenceConflict, UpdateNode(tx, n))
  1205  		return nil
  1206  	})
  1207  	assert.NoError(t, err)
  1208  
  1209  	// But using retrievedNode2 should work, since it has the latest
  1210  	// sequence information.
  1211  	retrievedNode2.Spec.Annotations.Name = "name3"
  1212  	err = s.Update(func(tx Tx) error {
  1213  		assert.NoError(t, UpdateNode(tx, retrievedNode2))
  1214  		return nil
  1215  	})
  1216  	assert.NoError(t, err)
  1217  }
  1218  
  1219  func TestTimestamps(t *testing.T) {
  1220  	s := NewMemoryStore(&testutils.MockProposer{})
  1221  	assert.NotNil(t, s)
  1222  
  1223  	var (
  1224  		retrievedNode *api.Node
  1225  		updatedNode   *api.Node
  1226  	)
  1227  
  1228  	// Create one node
  1229  	n := &api.Node{
  1230  		ID: "id1",
  1231  		Spec: api.NodeSpec{
  1232  			Annotations: api.Annotations{
  1233  				Name: "name1",
  1234  			},
  1235  		},
  1236  	}
  1237  	err := s.Update(func(tx Tx) error {
  1238  		assert.NoError(t, CreateNode(tx, n))
  1239  		return nil
  1240  	})
  1241  	assert.NoError(t, err)
  1242  
  1243  	// Make sure our local copy got updated.
  1244  	assert.NotZero(t, n.Meta.CreatedAt)
  1245  	assert.NotZero(t, n.Meta.UpdatedAt)
  1246  	// Since this is a new node, CreatedAt should equal UpdatedAt.
  1247  	assert.Equal(t, n.Meta.CreatedAt, n.Meta.UpdatedAt)
  1248  
  1249  	// Fetch the node from the store and make sure timestamps match.
  1250  	s.View(func(tx ReadTx) {
  1251  		retrievedNode = GetNode(tx, n.ID)
  1252  	})
  1253  	assert.Equal(t, retrievedNode.Meta.CreatedAt, n.Meta.CreatedAt)
  1254  	assert.Equal(t, retrievedNode.Meta.UpdatedAt, n.Meta.UpdatedAt)
  1255  
  1256  	// Make an update.
  1257  	retrievedNode.Spec.Annotations.Name = "name2"
  1258  	err = s.Update(func(tx Tx) error {
  1259  		assert.NoError(t, UpdateNode(tx, retrievedNode))
  1260  		updatedNode = GetNode(tx, n.ID)
  1261  		return nil
  1262  	})
  1263  	assert.NoError(t, err)
  1264  
  1265  	// Ensure `CreatedAt` is the same after the update and `UpdatedAt` got updated.
  1266  	assert.Equal(t, updatedNode.Meta.CreatedAt, n.Meta.CreatedAt)
  1267  	assert.NotEqual(t, updatedNode.Meta.CreatedAt, updatedNode.Meta.UpdatedAt)
  1268  }
  1269  
  1270  func TestBatch(t *testing.T) {
  1271  	s := NewMemoryStore(&testutils.MockProposer{})
  1272  	assert.NotNil(t, s)
  1273  
  1274  	watch, cancel := s.WatchQueue().Watch()
  1275  	defer cancel()
  1276  
  1277  	// Create 405 nodes. Should get split across 3 transactions.
  1278  	err := s.Batch(func(batch *Batch) error {
  1279  		for i := 0; i != 2*MaxChangesPerTransaction+5; i++ {
  1280  			n := &api.Node{
  1281  				ID: "id" + strconv.Itoa(i),
  1282  				Spec: api.NodeSpec{
  1283  					Annotations: api.Annotations{
  1284  						Name: "name" + strconv.Itoa(i),
  1285  					},
  1286  				},
  1287  			}
  1288  
  1289  			batch.Update(func(tx Tx) error {
  1290  				assert.NoError(t, CreateNode(tx, n))
  1291  				return nil
  1292  			})
  1293  		}
  1294  
  1295  		return nil
  1296  	})
  1297  	assert.NoError(t, err)
  1298  
  1299  	for i := 0; i != MaxChangesPerTransaction; i++ {
  1300  		event := <-watch
  1301  		if _, ok := event.(api.EventCreateNode); !ok {
  1302  			t.Fatalf("expected EventCreateNode; got %#v", event)
  1303  		}
  1304  	}
  1305  	event := <-watch
  1306  	if _, ok := event.(state.EventCommit); !ok {
  1307  		t.Fatalf("expected EventCommit; got %#v", event)
  1308  	}
  1309  	for i := 0; i != MaxChangesPerTransaction; i++ {
  1310  		event := <-watch
  1311  		if _, ok := event.(api.EventCreateNode); !ok {
  1312  			t.Fatalf("expected EventCreateNode; got %#v", event)
  1313  		}
  1314  	}
  1315  	event = <-watch
  1316  	if _, ok := event.(state.EventCommit); !ok {
  1317  		t.Fatalf("expected EventCommit; got %#v", event)
  1318  	}
  1319  	for i := 0; i != 5; i++ {
  1320  		event := <-watch
  1321  		if _, ok := event.(api.EventCreateNode); !ok {
  1322  			t.Fatalf("expected EventCreateNode; got %#v", event)
  1323  		}
  1324  	}
  1325  	event = <-watch
  1326  	if _, ok := event.(state.EventCommit); !ok {
  1327  		t.Fatalf("expected EventCommit; got %#v", event)
  1328  	}
  1329  }
  1330  
  1331  func TestBatchFailure(t *testing.T) {
  1332  	s := NewMemoryStore(&testutils.MockProposer{})
  1333  	assert.NotNil(t, s)
  1334  
  1335  	watch, cancel := s.WatchQueue().Watch()
  1336  	defer cancel()
  1337  
  1338  	// Return an error partway through a transaction.
  1339  	err := s.Batch(func(batch *Batch) error {
  1340  		for i := 0; ; i++ {
  1341  			n := &api.Node{
  1342  				ID: "id" + strconv.Itoa(i),
  1343  				Spec: api.NodeSpec{
  1344  					Annotations: api.Annotations{
  1345  						Name: "name" + strconv.Itoa(i),
  1346  					},
  1347  				},
  1348  			}
  1349  
  1350  			batch.Update(func(tx Tx) error {
  1351  				assert.NoError(t, CreateNode(tx, n))
  1352  				return nil
  1353  			})
  1354  			if i == MaxChangesPerTransaction+8 {
  1355  				return errors.New("failing the current tx")
  1356  			}
  1357  		}
  1358  	})
  1359  	assert.Error(t, err)
  1360  
  1361  	for i := 0; i != MaxChangesPerTransaction; i++ {
  1362  		event := <-watch
  1363  		if _, ok := event.(api.EventCreateNode); !ok {
  1364  			t.Fatalf("expected EventCreateNode; got %#v", event)
  1365  		}
  1366  	}
  1367  	event := <-watch
  1368  	if _, ok := event.(state.EventCommit); !ok {
  1369  		t.Fatalf("expected EventCommit; got %#v", event)
  1370  	}
  1371  
  1372  	// Shouldn't be anything after the first transaction
  1373  	select {
  1374  	case <-watch:
  1375  		t.Fatal("unexpected additional events")
  1376  	case <-time.After(50 * time.Millisecond):
  1377  	}
  1378  }
  1379  
  1380  func TestStoreSaveRestore(t *testing.T) {
  1381  	s1 := NewMemoryStore(nil)
  1382  	assert.NotNil(t, s1)
  1383  
  1384  	setupTestStore(t, s1)
  1385  
  1386  	var snapshot *api.StoreSnapshot
  1387  	s1.View(func(tx ReadTx) {
  1388  		var err error
  1389  		snapshot, err = s1.Save(tx)
  1390  		assert.NoError(t, err)
  1391  	})
  1392  
  1393  	s2 := NewMemoryStore(nil)
  1394  	assert.NotNil(t, s2)
  1395  	// setup s2 with the first element of each of the object sets (which should be
  1396  	// updated on restore), as well as one extraneous object (which should be deleted
  1397  	// on restore).  We also want to bump the version on all the ones that will be
  1398  	// updated just to make sure that restoration works.
  1399  	version := api.Version{Index: 100}
  1400  	c := clusterSet[0].Copy()
  1401  	c.Meta.Version = version
  1402  	n := nodeSet[0].Copy()
  1403  	n.Meta.Version = version
  1404  	s := serviceSet[0].Copy()
  1405  	s.Meta.Version = version
  1406  	task := taskSet[0].Copy()
  1407  	task.Meta.Version = version
  1408  	nw := networkSet[0].Copy()
  1409  	nw.Meta.Version = version
  1410  	cf := configSet[0].Copy()
  1411  	cf.Meta.Version = version
  1412  	sk := secretSet[0].Copy()
  1413  	sk.Meta.Version = version
  1414  	ext := extensionSet[0].Copy()
  1415  	ext.Meta.Version = version
  1416  	r := resourceSet[0].Copy()
  1417  	r.Meta.Version = version
  1418  	populateTestStore(t, s2,
  1419  		append(altClusterSet, c),
  1420  		append(altNodeSet, n),
  1421  		append(altServiceSet, s),
  1422  		append(altTaskSet, task),
  1423  		append(altNetworkSet, nw),
  1424  		append(altConfigSet, cf),
  1425  		append(altSecretSet, sk),
  1426  		append(altExtensionSet, ext),
  1427  		append(altResourceSet, r),
  1428  	)
  1429  
  1430  	watcher, cancel, err := ViewAndWatch(s2, func(ReadTx) error {
  1431  		return nil
  1432  	})
  1433  	assert.NoError(t, err)
  1434  	defer cancel()
  1435  
  1436  	err = s2.Restore(snapshot)
  1437  	assert.NoError(t, err)
  1438  
  1439  	// s2 should end up looking just like s1
  1440  	s2.View(func(tx ReadTx) {
  1441  		allClusters, err := FindClusters(tx, All)
  1442  		assert.NoError(t, err)
  1443  		assert.Len(t, allClusters, len(clusterSet))
  1444  		for i := range allClusters {
  1445  			assert.Equal(t, allClusters[i], clusterSet[i])
  1446  		}
  1447  
  1448  		allTasks, err := FindTasks(tx, All)
  1449  		assert.NoError(t, err)
  1450  		assert.Len(t, allTasks, len(taskSet))
  1451  		for i := range allTasks {
  1452  			assert.Equal(t, allTasks[i], taskSet[i])
  1453  		}
  1454  
  1455  		allNodes, err := FindNodes(tx, All)
  1456  		assert.NoError(t, err)
  1457  		assert.Len(t, allNodes, len(nodeSet))
  1458  		for i := range allNodes {
  1459  			assert.Equal(t, allNodes[i], nodeSet[i])
  1460  		}
  1461  
  1462  		allNetworks, err := FindNetworks(tx, All)
  1463  		assert.NoError(t, err)
  1464  		assert.Len(t, allNetworks, len(networkSet))
  1465  		for i := range allNetworks {
  1466  			assert.Equal(t, allNetworks[i], networkSet[i])
  1467  		}
  1468  
  1469  		allServices, err := FindServices(tx, All)
  1470  		assert.NoError(t, err)
  1471  		assert.Len(t, allServices, len(serviceSet))
  1472  		for i := range allServices {
  1473  			assert.Equal(t, allServices[i], serviceSet[i])
  1474  		}
  1475  
  1476  		allConfigs, err := FindConfigs(tx, All)
  1477  		assert.NoError(t, err)
  1478  		assert.Len(t, allConfigs, len(configSet))
  1479  		for i := range allConfigs {
  1480  			assert.Equal(t, allConfigs[i], configSet[i])
  1481  		}
  1482  
  1483  		allSecrets, err := FindSecrets(tx, All)
  1484  		assert.NoError(t, err)
  1485  		assert.Len(t, allSecrets, len(secretSet))
  1486  		for i := range allSecrets {
  1487  			assert.Equal(t, allSecrets[i], secretSet[i])
  1488  		}
  1489  
  1490  		allExtensions, err := FindExtensions(tx, All)
  1491  		assert.NoError(t, err)
  1492  		assert.Len(t, allExtensions, len(extensionSet))
  1493  		for i := range allExtensions {
  1494  			assert.Equal(t, allExtensions[i], extensionSet[i])
  1495  		}
  1496  
  1497  		allResources, err := FindResources(tx, All)
  1498  		assert.NoError(t, err)
  1499  		assert.Len(t, allResources, len(resourceSet))
  1500  		for i := range allResources {
  1501  			assert.Equal(t, allResources[i], resourceSet[i])
  1502  		}
  1503  	})
  1504  
  1505  	timeout := time.After(time.Second)
  1506  
  1507  	// make sure we have 1 update event, 2 create events, and 1 delete event for each
  1508  	// object type
  1509  	var (
  1510  		clusterUpdates, clusterCreates, clusterDeletes,
  1511  		nodeUpdates, nodeCreates, nodeDeletes,
  1512  		serviceUpdates, serviceCreates, serviceDeletes,
  1513  		taskUpdates, taskCreates, taskDeletes,
  1514  		networkUpdates, networkCreates, networkDeletes,
  1515  		configUpdates, configCreates, configDeletes,
  1516  		secretUpdates, secretCreates, secretDeletes,
  1517  		extensionUpdates, extensionCreates, extensionDeletes,
  1518  		resourceUpdates, resourceCreates, resourceDeletes []api.StoreObject
  1519  	)
  1520  
  1521  waitForAllEvents:
  1522  	for {
  1523  		var update events.Event
  1524  		select {
  1525  		case update = <-watcher:
  1526  		case <-timeout:
  1527  			assert.FailNow(t, "did not get all the events we were expecting after a snapshot was restored")
  1528  		}
  1529  
  1530  		switch e := update.(type) {
  1531  
  1532  		case api.EventUpdateCluster:
  1533  			clusterUpdates = append(clusterUpdates, e.Cluster)
  1534  		case api.EventCreateCluster:
  1535  			clusterCreates = append(clusterCreates, e.Cluster)
  1536  		case api.EventDeleteCluster:
  1537  			clusterDeletes = append(clusterDeletes, e.Cluster)
  1538  
  1539  		case api.EventUpdateNode:
  1540  			nodeUpdates = append(nodeUpdates, e.Node)
  1541  		case api.EventCreateNode:
  1542  			nodeCreates = append(nodeCreates, e.Node)
  1543  		case api.EventDeleteNode:
  1544  			nodeDeletes = append(nodeDeletes, e.Node)
  1545  
  1546  		case api.EventUpdateService:
  1547  			serviceUpdates = append(serviceUpdates, e.Service)
  1548  		case api.EventCreateService:
  1549  			serviceCreates = append(serviceCreates, e.Service)
  1550  		case api.EventDeleteService:
  1551  			serviceDeletes = append(serviceDeletes, e.Service)
  1552  
  1553  		case api.EventUpdateTask:
  1554  			taskUpdates = append(taskUpdates, e.Task)
  1555  		case api.EventCreateTask:
  1556  			taskCreates = append(taskCreates, e.Task)
  1557  		case api.EventDeleteTask:
  1558  			taskDeletes = append(taskDeletes, e.Task)
  1559  
  1560  		case api.EventUpdateNetwork:
  1561  			networkUpdates = append(networkUpdates, e.Network)
  1562  		case api.EventCreateNetwork:
  1563  			networkCreates = append(networkCreates, e.Network)
  1564  		case api.EventDeleteNetwork:
  1565  			networkDeletes = append(networkDeletes, e.Network)
  1566  
  1567  		case api.EventUpdateConfig:
  1568  			configUpdates = append(configUpdates, e.Config)
  1569  		case api.EventCreateConfig:
  1570  			configCreates = append(configCreates, e.Config)
  1571  		case api.EventDeleteConfig:
  1572  			configDeletes = append(configDeletes, e.Config)
  1573  
  1574  		case api.EventUpdateSecret:
  1575  			secretUpdates = append(secretUpdates, e.Secret)
  1576  		case api.EventCreateSecret:
  1577  			secretCreates = append(secretCreates, e.Secret)
  1578  		case api.EventDeleteSecret:
  1579  			secretDeletes = append(secretDeletes, e.Secret)
  1580  
  1581  		case api.EventUpdateExtension:
  1582  			extensionUpdates = append(extensionUpdates, e.Extension)
  1583  		case api.EventCreateExtension:
  1584  			extensionCreates = append(extensionCreates, e.Extension)
  1585  		case api.EventDeleteExtension:
  1586  			extensionDeletes = append(extensionDeletes, e.Extension)
  1587  
  1588  		case api.EventUpdateResource:
  1589  			resourceUpdates = append(resourceUpdates, e.Resource)
  1590  		case api.EventCreateResource:
  1591  			resourceCreates = append(resourceCreates, e.Resource)
  1592  		case api.EventDeleteResource:
  1593  			resourceDeletes = append(resourceDeletes, e.Resource)
  1594  		}
  1595  
  1596  		// wait until we have all the events we want
  1597  		for _, x := range [][]api.StoreObject{
  1598  			clusterUpdates, clusterDeletes,
  1599  			nodeUpdates, nodeDeletes,
  1600  			serviceUpdates, serviceDeletes,
  1601  			taskUpdates, taskDeletes,
  1602  			networkUpdates, networkDeletes,
  1603  			configUpdates, configDeletes,
  1604  			secretUpdates, secretDeletes,
  1605  			extensionUpdates, extensionDeletes,
  1606  			resourceUpdates, resourceDeletes,
  1607  		} {
  1608  			if len(x) < 1 {
  1609  				continue waitForAllEvents
  1610  			}
  1611  		}
  1612  
  1613  		for _, x := range [][]api.StoreObject{
  1614  			clusterCreates,
  1615  			nodeCreates,
  1616  			serviceCreates,
  1617  			taskCreates,
  1618  			networkCreates,
  1619  			configCreates,
  1620  			secretCreates,
  1621  			extensionCreates,
  1622  			resourceCreates,
  1623  		} {
  1624  			if len(x) < 2 {
  1625  				continue waitForAllEvents
  1626  			}
  1627  		}
  1628  		break
  1629  	}
  1630  
  1631  	assertHasSameIDs := func(changes []api.StoreObject, expected ...api.StoreObject) {
  1632  		assert.Equal(t, len(expected), len(changes))
  1633  		expectedIDs := make(map[string]struct{})
  1634  		for _, s := range expected {
  1635  			expectedIDs[s.GetID()] = struct{}{}
  1636  		}
  1637  		for _, s := range changes {
  1638  			_, ok := expectedIDs[s.GetID()]
  1639  			assert.True(t, ok)
  1640  		}
  1641  	}
  1642  
  1643  	assertHasSameIDs(clusterUpdates, clusterSet[0])
  1644  	assertHasSameIDs(clusterDeletes, altClusterSet[0])
  1645  	cantCastArrays := make([]api.StoreObject, len(clusterSet[1:]))
  1646  	for i, x := range clusterSet[1:] {
  1647  		cantCastArrays[i] = x
  1648  	}
  1649  	assertHasSameIDs(clusterCreates, cantCastArrays...)
  1650  
  1651  	assertHasSameIDs(nodeUpdates, nodeSet[0])
  1652  	assertHasSameIDs(nodeDeletes, altNodeSet[0])
  1653  	cantCastArrays = make([]api.StoreObject, len(nodeSet[1:]))
  1654  	for i, x := range nodeSet[1:] {
  1655  		cantCastArrays[i] = x
  1656  	}
  1657  	assertHasSameIDs(nodeCreates, cantCastArrays...)
  1658  
  1659  	assertHasSameIDs(serviceUpdates, serviceSet[0])
  1660  	assertHasSameIDs(serviceDeletes, altServiceSet[0])
  1661  	cantCastArrays = make([]api.StoreObject, len(serviceSet[1:]))
  1662  	for i, x := range serviceSet[1:] {
  1663  		cantCastArrays[i] = x
  1664  	}
  1665  	assertHasSameIDs(serviceCreates, cantCastArrays...)
  1666  
  1667  	assertHasSameIDs(taskUpdates, taskSet[0])
  1668  	assertHasSameIDs(taskDeletes, altTaskSet[0])
  1669  	cantCastArrays = make([]api.StoreObject, len(taskSet[1:]))
  1670  	for i, x := range taskSet[1:] {
  1671  		cantCastArrays[i] = x
  1672  	}
  1673  	assertHasSameIDs(taskCreates, cantCastArrays...)
  1674  
  1675  	assertHasSameIDs(networkUpdates, networkSet[0])
  1676  	assertHasSameIDs(networkDeletes, altNetworkSet[0])
  1677  	cantCastArrays = make([]api.StoreObject, len(networkSet[1:]))
  1678  	for i, x := range networkSet[1:] {
  1679  		cantCastArrays[i] = x
  1680  	}
  1681  	assertHasSameIDs(networkCreates, cantCastArrays...)
  1682  
  1683  	assertHasSameIDs(configUpdates, configSet[0])
  1684  	assertHasSameIDs(configDeletes, altConfigSet[0])
  1685  	cantCastArrays = make([]api.StoreObject, len(configSet[1:]))
  1686  	for i, x := range configSet[1:] {
  1687  		cantCastArrays[i] = x
  1688  	}
  1689  	assertHasSameIDs(configCreates, cantCastArrays...)
  1690  
  1691  	assertHasSameIDs(secretUpdates, secretSet[0])
  1692  	assertHasSameIDs(secretDeletes, altSecretSet[0])
  1693  	cantCastArrays = make([]api.StoreObject, len(secretSet[1:]))
  1694  	for i, x := range secretSet[1:] {
  1695  		cantCastArrays[i] = x
  1696  	}
  1697  	assertHasSameIDs(secretCreates, cantCastArrays...)
  1698  
  1699  	assertHasSameIDs(extensionUpdates, extensionSet[0])
  1700  	assertHasSameIDs(extensionDeletes, altExtensionSet[0])
  1701  	cantCastArrays = make([]api.StoreObject, len(extensionSet[1:]))
  1702  	for i, x := range extensionSet[1:] {
  1703  		cantCastArrays[i] = x
  1704  	}
  1705  	assertHasSameIDs(extensionCreates, cantCastArrays...)
  1706  
  1707  	assertHasSameIDs(resourceUpdates, resourceSet[0])
  1708  	assertHasSameIDs(resourceDeletes, altResourceSet[0])
  1709  	cantCastArrays = make([]api.StoreObject, len(resourceSet[1:]))
  1710  	for i, x := range resourceSet[1:] {
  1711  		cantCastArrays[i] = x
  1712  	}
  1713  	assertHasSameIDs(resourceCreates, cantCastArrays...)
  1714  }
  1715  
  1716  func TestWatchFrom(t *testing.T) {
  1717  	s := NewMemoryStore(&testutils.MockProposer{})
  1718  	assert.NotNil(t, s)
  1719  
  1720  	// Create a few nodes, 2 per transaction
  1721  	for i := 0; i != 5; i++ {
  1722  		err := s.Batch(func(batch *Batch) error {
  1723  			node := &api.Node{
  1724  				ID: "id" + strconv.Itoa(i),
  1725  				Spec: api.NodeSpec{
  1726  					Annotations: api.Annotations{
  1727  						Name: "name" + strconv.Itoa(i),
  1728  					},
  1729  				},
  1730  			}
  1731  
  1732  			service := &api.Service{
  1733  				ID: "id" + strconv.Itoa(i),
  1734  				Spec: api.ServiceSpec{
  1735  					Annotations: api.Annotations{
  1736  						Name: "name" + strconv.Itoa(i),
  1737  					},
  1738  				},
  1739  			}
  1740  
  1741  			batch.Update(func(tx Tx) error {
  1742  				assert.NoError(t, CreateNode(tx, node))
  1743  				return nil
  1744  			})
  1745  			batch.Update(func(tx Tx) error {
  1746  				assert.NoError(t, CreateService(tx, service))
  1747  				return nil
  1748  			})
  1749  			return nil
  1750  		})
  1751  		assert.NoError(t, err)
  1752  	}
  1753  
  1754  	// Try to watch from an invalid index
  1755  	_, _, err := WatchFrom(s, &api.Version{Index: 5000})
  1756  	assert.Error(t, err)
  1757  
  1758  	watch1, cancel1, err := WatchFrom(s, &api.Version{Index: 10}, api.EventCreateNode{}, state.EventCommit{})
  1759  	require.NoError(t, err)
  1760  	defer cancel1()
  1761  
  1762  	for i := 0; i != 2; i++ {
  1763  		select {
  1764  		case event := <-watch1:
  1765  			nodeEvent, ok := event.(api.EventCreateNode)
  1766  			if !ok {
  1767  				t.Fatal("wrong event type - expected node create")
  1768  			}
  1769  
  1770  			if i == 0 {
  1771  				assert.Equal(t, "id3", nodeEvent.Node.ID)
  1772  			} else {
  1773  				assert.Equal(t, "id4", nodeEvent.Node.ID)
  1774  			}
  1775  		case <-time.After(time.Second):
  1776  			t.Fatal("timed out waiting for event")
  1777  		}
  1778  		select {
  1779  		case event := <-watch1:
  1780  			if _, ok := event.(state.EventCommit); !ok {
  1781  				t.Fatal("wrong event type - expected commit")
  1782  			}
  1783  		case <-time.After(time.Second):
  1784  			t.Fatal("timed out waiting for event")
  1785  		}
  1786  	}
  1787  
  1788  	watch2, cancel2, err := WatchFrom(s, &api.Version{Index: 13}, api.EventCreateService{}, state.EventCommit{})
  1789  	require.NoError(t, err)
  1790  	defer cancel2()
  1791  
  1792  	select {
  1793  	case event := <-watch2:
  1794  		serviceEvent, ok := event.(api.EventCreateService)
  1795  		if !ok {
  1796  			t.Fatal("wrong event type - expected service create")
  1797  		}
  1798  		assert.Equal(t, "id4", serviceEvent.Service.ID)
  1799  	case <-time.After(time.Second):
  1800  		t.Fatal("timed out waiting for event")
  1801  	}
  1802  	select {
  1803  	case event := <-watch2:
  1804  		if _, ok := event.(state.EventCommit); !ok {
  1805  			t.Fatal("wrong event type - expected commit")
  1806  		}
  1807  	case <-time.After(time.Second):
  1808  		t.Fatal("timed out waiting for event")
  1809  	}
  1810  
  1811  	// Create some new objects and make sure they show up in the watches.
  1812  	assert.NoError(t, s.Update(func(tx Tx) error {
  1813  		node := &api.Node{
  1814  			ID: "newnode",
  1815  			Spec: api.NodeSpec{
  1816  				Annotations: api.Annotations{
  1817  					Name: "newnode",
  1818  				},
  1819  			},
  1820  		}
  1821  
  1822  		service := &api.Service{
  1823  			ID: "newservice",
  1824  			Spec: api.ServiceSpec{
  1825  				Annotations: api.Annotations{
  1826  					Name: "newservice",
  1827  				},
  1828  			},
  1829  		}
  1830  
  1831  		assert.NoError(t, CreateNode(tx, node))
  1832  		assert.NoError(t, CreateService(tx, service))
  1833  		return nil
  1834  	}))
  1835  
  1836  	select {
  1837  	case event := <-watch1:
  1838  		nodeEvent, ok := event.(api.EventCreateNode)
  1839  		if !ok {
  1840  			t.Fatalf("wrong event type - expected node create, got %T", event)
  1841  		}
  1842  		assert.Equal(t, "newnode", nodeEvent.Node.ID)
  1843  	case <-time.After(time.Second):
  1844  		t.Fatal("timed out waiting for event")
  1845  	}
  1846  	select {
  1847  	case event := <-watch1:
  1848  		if _, ok := event.(state.EventCommit); !ok {
  1849  			t.Fatal("wrong event type - expected commit")
  1850  		}
  1851  	case <-time.After(time.Second):
  1852  		t.Fatal("timed out waiting for event")
  1853  	}
  1854  
  1855  	select {
  1856  	case event := <-watch2:
  1857  		serviceEvent, ok := event.(api.EventCreateService)
  1858  		if !ok {
  1859  			t.Fatalf("wrong event type - expected service create, got %T", event)
  1860  		}
  1861  		assert.Equal(t, "newservice", serviceEvent.Service.ID)
  1862  	case <-time.After(time.Second):
  1863  		t.Fatal("timed out waiting for event")
  1864  	}
  1865  	select {
  1866  	case event := <-watch2:
  1867  		if _, ok := event.(state.EventCommit); !ok {
  1868  			t.Fatal("wrong event type - expected commit")
  1869  		}
  1870  	case <-time.After(time.Second):
  1871  		t.Fatal("timed out waiting for event")
  1872  	}
  1873  
  1874  	assert.NoError(t, s.Update(func(tx Tx) error {
  1875  		node := &api.Node{
  1876  			ID: "newnode2",
  1877  			Spec: api.NodeSpec{
  1878  				Annotations: api.Annotations{
  1879  					Name: "newnode2",
  1880  				},
  1881  			},
  1882  		}
  1883  
  1884  		assert.NoError(t, CreateNode(tx, node))
  1885  		return nil
  1886  	}))
  1887  
  1888  	select {
  1889  	case event := <-watch1:
  1890  		nodeEvent, ok := event.(api.EventCreateNode)
  1891  		if !ok {
  1892  			t.Fatalf("wrong event type - expected node create, got %T", event)
  1893  		}
  1894  		assert.Equal(t, "newnode2", nodeEvent.Node.ID)
  1895  	case <-time.After(time.Second):
  1896  		t.Fatal("timed out waiting for event")
  1897  	}
  1898  	select {
  1899  	case event := <-watch1:
  1900  		if _, ok := event.(state.EventCommit); !ok {
  1901  			t.Fatal("wrong event type - expected commit")
  1902  		}
  1903  	case <-time.After(time.Second):
  1904  		t.Fatal("timed out waiting for event")
  1905  	}
  1906  
  1907  	select {
  1908  	case event := <-watch2:
  1909  		if _, ok := event.(state.EventCommit); !ok {
  1910  			t.Fatal("wrong event type - expected commit")
  1911  		}
  1912  	case <-time.After(time.Second):
  1913  		t.Fatal("timed out waiting for event")
  1914  	}
  1915  }
  1916  
  1917  const benchmarkNumNodes = 10000
  1918  
  1919  func setupNodes(b *testing.B, n int) (*MemoryStore, []string) {
  1920  	s := NewMemoryStore(nil)
  1921  
  1922  	nodeIDs := make([]string, n)
  1923  
  1924  	for i := 0; i < n; i++ {
  1925  		nodeIDs[i] = identity.NewID()
  1926  	}
  1927  
  1928  	b.ResetTimer()
  1929  
  1930  	_ = s.Update(func(tx1 Tx) error {
  1931  		for i := 0; i < n; i++ {
  1932  			_ = CreateNode(tx1, &api.Node{
  1933  				ID: nodeIDs[i],
  1934  				Spec: api.NodeSpec{
  1935  					Annotations: api.Annotations{
  1936  						Name: "name" + strconv.Itoa(i),
  1937  					},
  1938  				},
  1939  			})
  1940  		}
  1941  		return nil
  1942  	})
  1943  
  1944  	return s, nodeIDs
  1945  }
  1946  
  1947  func BenchmarkCreateNode(b *testing.B) {
  1948  	setupNodes(b, b.N)
  1949  }
  1950  
  1951  func BenchmarkUpdateNode(b *testing.B) {
  1952  	s, nodeIDs := setupNodes(b, benchmarkNumNodes)
  1953  	b.ResetTimer()
  1954  	_ = s.Update(func(tx1 Tx) error {
  1955  		for i := 0; i < b.N; i++ {
  1956  			_ = UpdateNode(tx1, &api.Node{
  1957  				ID: nodeIDs[i%benchmarkNumNodes],
  1958  				Spec: api.NodeSpec{
  1959  					Annotations: api.Annotations{
  1960  						Name: nodeIDs[i%benchmarkNumNodes] + "_" + strconv.Itoa(i),
  1961  					},
  1962  				},
  1963  			})
  1964  		}
  1965  		return nil
  1966  	})
  1967  }
  1968  
  1969  func BenchmarkUpdateNodeTransaction(b *testing.B) {
  1970  	s, nodeIDs := setupNodes(b, benchmarkNumNodes)
  1971  	b.ResetTimer()
  1972  	for i := 0; i < b.N; i++ {
  1973  		_ = s.Update(func(tx1 Tx) error {
  1974  			_ = UpdateNode(tx1, &api.Node{
  1975  				ID: nodeIDs[i%benchmarkNumNodes],
  1976  				Spec: api.NodeSpec{
  1977  					Annotations: api.Annotations{
  1978  						Name: nodeIDs[i%benchmarkNumNodes] + "_" + strconv.Itoa(i),
  1979  					},
  1980  				},
  1981  			})
  1982  			return nil
  1983  		})
  1984  	}
  1985  }
  1986  
  1987  func BenchmarkDeleteNodeTransaction(b *testing.B) {
  1988  	s, nodeIDs := setupNodes(b, benchmarkNumNodes)
  1989  	b.ResetTimer()
  1990  	for i := 0; i < b.N; i++ {
  1991  		_ = s.Update(func(tx1 Tx) error {
  1992  			_ = DeleteNode(tx1, nodeIDs[0])
  1993  			// Don't actually commit deletions, so we can delete
  1994  			// things repeatedly to satisfy the benchmark structure.
  1995  			return errors.New("don't commit this")
  1996  		})
  1997  	}
  1998  }
  1999  
  2000  func BenchmarkGetNode(b *testing.B) {
  2001  	s, nodeIDs := setupNodes(b, benchmarkNumNodes)
  2002  	b.ResetTimer()
  2003  	s.View(func(tx1 ReadTx) {
  2004  		for i := 0; i < b.N; i++ {
  2005  			_ = GetNode(tx1, nodeIDs[i%benchmarkNumNodes])
  2006  		}
  2007  	})
  2008  }
  2009  
  2010  func BenchmarkFindAllNodes(b *testing.B) {
  2011  	s, _ := setupNodes(b, benchmarkNumNodes)
  2012  	b.ResetTimer()
  2013  	s.View(func(tx1 ReadTx) {
  2014  		for i := 0; i < b.N; i++ {
  2015  			_, _ = FindNodes(tx1, All)
  2016  		}
  2017  	})
  2018  }
  2019  
  2020  func BenchmarkFindNodeByName(b *testing.B) {
  2021  	s, _ := setupNodes(b, benchmarkNumNodes)
  2022  	b.ResetTimer()
  2023  	s.View(func(tx1 ReadTx) {
  2024  		for i := 0; i < b.N; i++ {
  2025  			_, _ = FindNodes(tx1, ByName("name"+strconv.Itoa(i)))
  2026  		}
  2027  	})
  2028  }
  2029  
  2030  func BenchmarkNodeConcurrency(b *testing.B) {
  2031  	s, nodeIDs := setupNodes(b, benchmarkNumNodes)
  2032  	b.ResetTimer()
  2033  
  2034  	// Run 5 writer goroutines and 5 reader goroutines
  2035  	var wg sync.WaitGroup
  2036  	for c := 0; c != 5; c++ {
  2037  		wg.Add(1)
  2038  		go func(c int) {
  2039  			defer wg.Done()
  2040  			for i := 0; i < b.N; i++ {
  2041  				_ = s.Update(func(tx1 Tx) error {
  2042  					_ = UpdateNode(tx1, &api.Node{
  2043  						ID: nodeIDs[i%benchmarkNumNodes],
  2044  						Spec: api.NodeSpec{
  2045  							Annotations: api.Annotations{
  2046  								Name: nodeIDs[i%benchmarkNumNodes] + "_" + strconv.Itoa(c) + "_" + strconv.Itoa(i),
  2047  							},
  2048  						},
  2049  					})
  2050  					return nil
  2051  				})
  2052  			}
  2053  		}(c)
  2054  	}
  2055  
  2056  	for c := 0; c != 5; c++ {
  2057  		wg.Add(1)
  2058  		go func() {
  2059  			defer wg.Done()
  2060  			s.View(func(tx1 ReadTx) {
  2061  				for i := 0; i < b.N; i++ {
  2062  					_ = GetNode(tx1, nodeIDs[i%benchmarkNumNodes])
  2063  				}
  2064  			})
  2065  		}()
  2066  	}
  2067  
  2068  	wg.Wait()
  2069  }