func updateNodeAvailability(t *testing.T, s *store.MemoryStore, node *api.Node, avail api.NodeSpec_Availability) {
node.Spec.Availability = avail
s.Update(func(tx store.Tx) error {
assert.NoError(t, store.UpdateNode(tx, node))
return nil
})
}
// getRunnableAndDeadSlots returns two maps of slots. The first contains slots
// that have at least one task with a desired state above NEW and lesser or
// equal to RUNNING. The second is for slots that only contain tasks with a
// desired state above RUNNING.
func getRunnableAndDeadSlots(s *store.MemoryStore, serviceID string) (map[uint64]slot, map[uint64]slot, error) {
var (
tasks []*api.Task
err error
)
s.View(func(tx store.ReadTx) {
tasks, err = store.FindTasks(tx, store.ByServiceID(serviceID))
})
if err != nil {
return nil, nil, err
}
runningSlots := make(map[uint64]slot)
for _, t := range tasks {
if t.DesiredState <= api.TaskStateRunning {
runningSlots[t.Slot] = append(runningSlots[t.Slot], t)
}
}
deadSlots := make(map[uint64]slot)
for _, t := range tasks {
if _, exists := runningSlots[t.Slot]; !exists {
deadSlots[t.Slot] = append(deadSlots[t.Slot], t)
}
}
return runningSlots, deadSlots, nil
}
func createNode(s *store.MemoryStore, nodeID, role string, csr, cert []byte) error {
apiRole, _ := ca.FormatRole(role)
err := s.Update(func(tx store.Tx) error {
node := &api.Node{
ID: nodeID,
Certificate: api.Certificate{
CSR: csr,
CN: nodeID,
Role: apiRole,
Status: api.IssuanceStatus{
State: api.IssuanceStateIssued,
},
Certificate: cert,
},
Spec: api.NodeSpec{
Role: apiRole,
Membership: api.NodeMembershipAccepted,
},
}
return store.CreateNode(tx, node)
})
return err
}
func deleteServiceTasks(ctx context.Context, s *store.MemoryStore, service *api.Service) {
var (
tasks []*api.Task
err error
)
s.View(func(tx store.ReadTx) {
tasks, err = store.FindTasks(tx, store.ByServiceID(service.ID))
})
if err != nil {
log.G(ctx).WithError(err).Errorf("failed to list tasks")
return
}
_, err = s.Batch(func(batch *store.Batch) error {
for _, t := range tasks {
err := batch.Update(func(tx store.Tx) error {
if err := store.DeleteTask(tx, t.ID); err != nil {
log.G(ctx).WithError(err).Errorf("failed to delete task")
}
return nil
})
if err != nil {
return err
}
}
return nil
})
if err != nil {
log.G(ctx).WithError(err).Errorf("task search transaction failed")
}
}
func createCluster(t *testing.T, s *store.MemoryStore, id, name string) *api.Cluster {
spec := createClusterSpec(name)
cluster := &api.Cluster{
ID: id,
Spec: *spec,
}
assert.NoError(t, s.Update(func(tx store.Tx) error {
return store.CreateCluster(tx, cluster)
}))
return cluster
}
func createClusterObject(t *testing.T, s *store.MemoryStore, acceptancePolicy api.AcceptancePolicy) {
assert.NoError(t, s.Update(func(tx store.Tx) error {
store.CreateCluster(tx, &api.Cluster{
ID: identity.NewID(),
Spec: api.ClusterSpec{
Annotations: api.Annotations{
Name: store.DefaultClusterName,
},
AcceptancePolicy: acceptancePolicy,
},
})
return nil
}))
}
func createClusterObject(t *testing.T, s *store.MemoryStore, clusterID string, acceptancePolicy api.AcceptancePolicy, externalCAs ...*api.ExternalCA) {
assert.NoError(t, s.Update(func(tx store.Tx) error {
store.CreateCluster(tx, &api.Cluster{
ID: clusterID,
Spec: api.ClusterSpec{
Annotations: api.Annotations{
Name: store.DefaultClusterName,
},
AcceptancePolicy: acceptancePolicy,
CAConfig: api.CAConfig{
ExternalCAs: externalCAs,
},
},
})
return nil
}))
}
func isValidEndpoint(t assert.TestingT, s *store.MemoryStore, task *api.Task) bool {
if task.ServiceID != "" {
var service *api.Service
s.View(func(tx store.ReadTx) {
service = store.GetService(tx, task.ServiceID)
})
if service == nil {
return true
}
return assert.Equal(t, service.Endpoint, task.Endpoint)
}
return true
}
// NewUpdater creates a new Updater.
func NewUpdater(store *store.MemoryStore, restartSupervisor *RestartSupervisor) *Updater {
return &Updater{
store: store,
watchQueue: store.WatchQueue(),
restarts: restartSupervisor,
stopChan: make(chan struct{}),
doneChan: make(chan struct{}),
}
}
func getRunnableServiceTasks(t *testing.T, s *store.MemoryStore, service *api.Service) []*api.Task {
var (
err error
tasks []*api.Task
)
s.View(func(tx store.ReadTx) {
tasks, err = store.FindTasks(tx, store.ByServiceID(service.ID))
})
assert.NoError(t, err)
runnable := []*api.Task{}
for _, task := range tasks {
if task.DesiredState == api.TaskStateRunning {
runnable = append(runnable, task)
}
}
return runnable
}
// NewUpdater creates a new Updater.
func NewUpdater(store *store.MemoryStore, restartSupervisor *RestartSupervisor, cluster *api.Cluster, newService *api.Service) *Updater {
return &Updater{
store: store,
watchQueue: store.WatchQueue(),
restarts: restartSupervisor,
cluster: cluster.Copy(),
newService: newService.Copy(),
stopChan: make(chan struct{}),
doneChan: make(chan struct{}),
}
}
// NewTaskReaper creates a new TaskReaper.
func NewTaskReaper(store *store.MemoryStore) *TaskReaper {
watcher, cancel := state.Watch(store.WatchQueue(), state.EventCreateTask{}, state.EventUpdateCluster{})
return &TaskReaper{
store: store,
watcher: watcher,
cancelWatch: cancel,
dirty: make(map[instanceTuple]struct{}),
stopChan: make(chan struct{}),
doneChan: make(chan struct{}),
}
}
func createClusterObject(t *testing.T, s *store.MemoryStore, clusterID, workerToken, managerToken string, externalCAs ...*api.ExternalCA) {
assert.NoError(t, s.Update(func(tx store.Tx) error {
store.CreateCluster(tx, &api.Cluster{
ID: clusterID,
Spec: api.ClusterSpec{
Annotations: api.Annotations{
Name: store.DefaultClusterName,
},
CAConfig: api.CAConfig{
ExternalCAs: externalCAs,
},
},
RootCA: api.RootCA{
JoinTokens: api.JoinTokens{
Worker: workerToken,
Manager: managerToken,
},
},
})
return nil
}))
}
// getRunnableSlots returns a map of slots that have at least one task with
// a desired state above NEW and lesser or equal to RUNNING.
func getRunnableSlots(s *store.MemoryStore, serviceID string) (map[uint64]slot, error) {
var (
tasks []*api.Task
err error
)
s.View(func(tx store.ReadTx) {
tasks, err = store.FindTasks(tx, store.ByServiceID(serviceID))
})
if err != nil {
return nil, err
}
runningSlots := make(map[uint64]slot)
for _, t := range tasks {
// Technically the check below could just be
// t.DesiredState <= api.TaskStateRunning, but ignoring tasks
// with DesiredState == NEW simplifies the drainer unit tests.
if t.DesiredState > api.TaskStateNew && t.DesiredState <= api.TaskStateRunning {
runningSlots[t.Slot] = append(runningSlots[t.Slot], t)
}
}
return runningSlots, nil
}
func deleteTask(t *testing.T, s *store.MemoryStore, task *api.Task) {
s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteTask(tx, task.ID))
return nil
})
}
func deleteNode(t *testing.T, s *store.MemoryStore, node *api.Node) {
s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteNode(tx, node.ID))
return nil
})
}
func deleteService(t *testing.T, s *store.MemoryStore, service *api.Service) {
s.Update(func(tx store.Tx) error {
assert.NoError(t, store.DeleteService(tx, service.ID))
return nil
})
}
