reboot: implement data struct and storage API

Signed-off-by: Daichi Sakaue <daichi-sakaue@cybozu.co.jp>
Signed-off-by: morimoto-cybozu <kenji_morimoto@cybozu.co.jp>

Author: yokaze

docs/reboot.md

@@ -23,11 +23,11 @@ Data schema
 
 ### `RebootQueueEntry`
 
-| Name     | Type     | Description                                |
-| -------- | -------- | ------------------------------------------ |
-| `index`  | uint64   | Index number of entry.                     |
-| `nodes`  | []string | A list of IP addresses of nodes to reboot. |
-| `status` | string   | One of `queued`, `rebooting`, `cancelled`. |
+| Name     | Type     | Description                                   |
+| -------- | -------- | --------------------------------------------- |
+| `index`  | string   | Index number of entry, formatted as a string. |
+| `nodes`  | []string | A list of IP addresses of nodes to reboot.    |
+| `status` | string   | One of `queued`, `rebooting`, `cancelled`.    |
 
 
 Detailed behavior

reboot.go

@@ -0,0 +1,32 @@
+package cke
+
+// RebootStatus is status of reboot operation
+type RebootStatus string
+
+// Reboot statuses
+const (
+	RebootStatusQueued    = RebootStatus("queued")
+	RebootStatusRebooting = RebootStatus("rebooting")
+	RebootStatusCancelled = RebootStatus("cancelled")
+)
+
+// RebootQueueEntry represents a queue entry of reboot operation
+type RebootQueueEntry struct {
+	Index  int64        `json:"index,string"`
+	Nodes  []string     `json:"nodes"`
+	Status RebootStatus `json:"status"`
+}
+
+// NewRebootQueueEntry creates new `RebootQueueEntry`.
+// `Index` will be supplied in registration.
+func NewRebootQueueEntry(nodes []string) *RebootQueueEntry {
+	return &RebootQueueEntry{
+		Nodes:  nodes,
+		Status: RebootStatusQueued,
+	}
+}
+
+// Cancel cancels the entry.
+func (r *RebootQueueEntry) Cancel() {
+	r.Status = RebootStatusCancelled
+}

storage.go

@@ -30,6 +30,8 @@ const (
 	KeyClusterRevision       = "cluster-revision"
 	KeyConstraints           = "constraints"
 	KeyLeader                = "leader/"
+	KeyRebootsPrefix         = "reboots/data/"
+	KeyRebootsWriteIndex     = "reboots/write-index"
 	KeyRecords               = "records/"
 	KeyRecordID              = "records"
 	KeyResourcePrefix        = "resource/"
@@ -685,6 +687,179 @@ func (s Storage) GetSabakanURL(ctx context.Context) (string, error) {
 	return s.getStringValue(ctx, KeySabakanURL)
 }
 
+func rebootsEntryKey(index int64) string {
+	return fmt.Sprintf("%s%016x", KeyRebootsPrefix, index)
+}
+
+// RegisterRebootsEntry enqueues a reboot queue entry to the reboot queue.
+// "Index" of the entry is retrieved and updated in this method. The given value is ignored.
+func (s Storage) RegisterRebootsEntry(ctx context.Context, r *RebootQueueEntry) error {
+RETRY:
+	var writeIndex, writeIndexRev int64
+	resp, err := s.Get(ctx, KeyRebootsWriteIndex)
+	if err != nil {
+		return err
+	}
+	if resp.Count != 0 {
+		value, err := strconv.ParseInt(string(resp.Kvs[0].Value), 10, 64)
+		if err != nil {
+			return err
+		}
+		writeIndex = value
+		writeIndexRev = resp.Kvs[0].ModRevision
+	}
+
+	r.Index = writeIndex
+	data, err := json.Marshal(r)
+	if err != nil {
+		return err
+	}
+
+	nextWriteIndex := strconv.FormatInt(writeIndex+1, 10)
+	txnResp, err := s.Txn(ctx).
+		If(
+			clientv3.Compare(clientv3.ModRevision(KeyRebootsWriteIndex), "=", writeIndexRev),
+		).
+		Then(
+			clientv3.OpPut(rebootsEntryKey(writeIndex), string(data)),
+			clientv3.OpPut(KeyRebootsWriteIndex, nextWriteIndex),
+		).
+		Commit()
+	if err != nil {
+		return err
+	}
+	if !txnResp.Succeeded {
+		goto RETRY
+	}
+
+	return nil
+}
+
+// UpdateRebootsEntry updates existing reboot queue entry.
+// It always overwrites the contents with a CAS loop.
+// If the entry is not found in the reboot queue, this returns ErrNotFound.
+func (s Storage) UpdateRebootsEntry(ctx context.Context, r *RebootQueueEntry) error {
+	key := rebootsEntryKey(r.Index)
+	data, err := json.Marshal(r)
+	if err != nil {
+		return err
+	}
+
+RETRY:
+	resp, err := s.Get(ctx, key)
+	if err != nil {
+		return err
+	}
+	if resp.Count == 0 {
+		return ErrNotFound
+	}
+
+	rev := resp.Kvs[0].ModRevision
+	txnResp, err := s.Txn(ctx).
+		If(
+			clientv3.Compare(clientv3.ModRevision(key), "=", rev),
+		).
+		Then(
+			clientv3.OpPut(key, string(data)),
+		).
+		Commit()
+	if err != nil {
+		return err
+	}
+	if !txnResp.Succeeded {
+		goto RETRY
+	}
+
+	return nil
+}
+
+// GetRebootsEntry loads the entry specified by the index from the reboot queue.
+// If the pointed entry is not found, this returns ErrNotFound.
+func (s Storage) GetRebootsEntry(ctx context.Context, index int64) (*RebootQueueEntry, error) {
+	resp, err := s.Get(ctx, rebootsEntryKey(index))
+	if err != nil {
+		return nil, err
+	}
+
+	if len(resp.Kvs) == 0 {
+		return nil, ErrNotFound
+	}
+
+	r := new(RebootQueueEntry)
+	err = json.Unmarshal(resp.Kvs[0].Value, r)
+	if err != nil {
+		return nil, err
+	}
+
+	return r, nil
+}
+
+func (s Storage) getRebootsEntries(ctx context.Context, count int64) ([]*RebootQueueEntry, error) {
+	opts := []clientv3.OpOption{
+		clientv3.WithPrefix(),
+		clientv3.WithSort(clientv3.SortByKey, clientv3.SortAscend),
+	}
+	if count > 0 {
+		opts = append(opts, clientv3.WithLimit(count))
+	}
+	resp, err := s.Get(ctx, KeyRebootsPrefix, opts...)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(resp.Kvs) == 0 {
+		return nil, nil
+	}
+
+	reboots := make([]*RebootQueueEntry, len(resp.Kvs))
+	for i, kv := range resp.Kvs {
+		r := new(RebootQueueEntry)
+		err = json.Unmarshal(kv.Value, r)
+		if err != nil {
+			return nil, err
+		}
+		reboots[i] = r
+	}
+
+	return reboots, nil
+}
+
+// GetRebootsEntries loads the entries from the reboot queue.
+func (s Storage) GetRebootsEntries(ctx context.Context) ([]*RebootQueueEntry, error) {
+	return s.getRebootsEntries(ctx, 0)
+}
+
+// GetRebootsFrontEntry loads the front entry from the reboot queue.
+// If the queue is empty, this returns ErrNotFound.
+func (s Storage) GetRebootsFrontEntry(ctx context.Context) (*RebootQueueEntry, error) {
+	reboots, err := s.getRebootsEntries(ctx, 1)
+	if err != nil {
+		return nil, err
+	}
+
+	if len(reboots) == 0 {
+		return nil, ErrNotFound
+	}
+
+	return reboots[0], nil
+}
+
+// DeleteRebootsEntry deletes the entry specified by the index from the reboot queue.
+func (s Storage) DeleteRebootsEntry(ctx context.Context, leaderKey string, index int64) error {
+	resp, err := s.Txn(ctx).
+		If(clientv3util.KeyExists(leaderKey)).
+		Then(clientv3.OpDelete(rebootsEntryKey(index))).
+		Commit()
+	if err != nil {
+		return err
+	}
+	if !resp.Succeeded {
+		return ErrNoLeader
+	}
+
+	return nil
+}
+
 // SetStatus stores the server status.
 func (s Storage) SetStatus(ctx context.Context, lease clientv3.LeaseID, st *ServerStatus) error {
 	data, err := json.Marshal(st)

storage_test.go

@@ -627,6 +627,121 @@ func testStorageSabakan(t *testing.T) {
 	}
 }
 
+func testStorageReboot(t *testing.T) {
+	t.Parallel()
+
+	client := newEtcdClient(t)
+	defer client.Close()
+	storage := Storage{client}
+	ctx := context.Background()
+
+	s, err := concurrency.NewSession(client)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer s.Close()
+	e := concurrency.NewElection(s, KeyLeader)
+	err = e.Campaign(ctx, "test")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	leaderKey := e.Key()
+
+	_, err = storage.GetRebootsEntry(ctx, 0)
+	if err != ErrNotFound {
+		t.Error("unexptected error:", err)
+	}
+
+	_, err = storage.GetRebootsFrontEntry(ctx)
+	if err != ErrNotFound {
+		t.Error("unexptected error:", err)
+	}
+
+	ents, err := storage.GetRebootsEntries(ctx)
+	if err != nil {
+		t.Fatal("GetRebootsEntries failed:", err)
+	}
+	if len(ents) != 0 {
+		t.Error("Unknown entries:", ents)
+	}
+
+	err = storage.DeleteRebootsEntry(ctx, leaderKey, 0)
+	if err != nil {
+		t.Fatal("DeleteRebootsEntry failed:", err)
+	}
+
+	nodes := []string{
+		"1.2.3.4",
+		"5.6.7.8",
+	}
+	entry := NewRebootQueueEntry(nodes)
+	err = storage.RegisterRebootsEntry(ctx, entry)
+	if err != nil {
+		t.Fatal("RegisterRebootsEntry failed:", err)
+	}
+
+	nodes2 := []string{
+		"12.34.56.78",
+	}
+	entry2 := NewRebootQueueEntry(nodes2)
+	err = storage.RegisterRebootsEntry(ctx, entry2)
+	if err != nil {
+		t.Fatal("RegisterRebootsEntry failed:", err)
+	}
+
+	ent, err := storage.GetRebootsEntry(ctx, 1)
+	if err != nil {
+		t.Fatal("GetRebootsEntry failed:", err)
+	}
+	if !cmp.Equal(ent, entry2) {
+		t.Error("GetRebootsEntry returned unexpected result:", cmp.Diff(ent, entry2))
+	}
+
+	ent, err = storage.GetRebootsFrontEntry(ctx)
+	if err != nil {
+		t.Fatal("GetRebootsFrontEntry failed:", err)
+	}
+	if !cmp.Equal(ent, entry) {
+		t.Error("GetRebootsFrontEntry returned unexpected result:", cmp.Diff(ent, entry))
+	}
+
+	entries := []*RebootQueueEntry{entry, entry2}
+	ents, err = storage.GetRebootsEntries(ctx)
+	if err != nil {
+		t.Fatal("GetRebootsEntries failed:", err)
+	}
+	if !cmp.Equal(ents, entries) {
+		t.Error("GetRebootsEntries returned unexpected result:", cmp.Diff(ents, entries))
+	}
+
+	entry.Status = RebootStatusRebooting
+	err = storage.UpdateRebootsEntry(ctx, entry)
+	if err != nil {
+		t.Fatal("UpdateRebootsEntry failed:", err)
+	}
+	ent, err = storage.GetRebootsEntry(ctx, 0)
+	if err != nil {
+		t.Fatal("GetRebootsEntry failed:", err)
+	}
+	if !cmp.Equal(ent, entry) {
+		t.Error("GetRebootsFrontEntry returned unexpected result:", cmp.Diff(ent, entry))
+	}
+
+	err = storage.DeleteRebootsEntry(ctx, leaderKey, 0)
+	if err != nil {
+		t.Fatal("DeleteRebootsEntry failed:", err)
+	}
+	_, err = storage.GetRebootsEntry(ctx, 0)
+	if err != ErrNotFound {
+		t.Error("unexptected error:", err)
+	}
+	err = storage.UpdateRebootsEntry(ctx, entry)
+	if err == nil {
+		t.Error("UpdateRebootsEntry succeeded for deleted entry")
+	}
+}
+
 func testStatus(t *testing.T) {
 	t.Parallel()
 
@@ -677,5 +792,6 @@ func TestStorage(t *testing.T) {
 	t.Run("Maint", testStorageMaint)
 	t.Run("Resource", testStorageResource)
 	t.Run("Sabakan", testStorageSabakan)
+	t.Run("Reboot", testStorageReboot)
 	t.Run("Status", testStatus)
 }