API: introduce ability to watch a VM (#351)

* API: introduce ability to watch a VM

* Document ?watch=true for GET /vms/{name} in the OpenAPI specification

* WatchVM: ensure that goroutine is terminated on early return with error

* WatchVM: close channels on goroutine exit

* WatchVM: ensure that we wait for the goroutine after additional barriers

* WatchVM: ignore unexpected keys instead of throwing an error

* WatchVM: perform context-aware writes to a bounded channel

* WatchVM: don't forget to close errCh on goroutine exit too

* WatchVM: don't close readyCh in goroutine to avoid ambiguity

* WatchVM: filter out spurious KVs that signify VM deletion

internal/controller/store/badger/badger_vm.go

@@ -3,9 +3,10 @@ package badger
 
 import (
 	"encoding/json"
+	"path"
+
 	"github.com/cirruslabs/orchard/pkg/resource/v1"
 	"github.com/dgraph-io/badger/v3"
-	"path"
 )
 
 const SpaceVMs = "/vms"
@@ -38,6 +39,8 @@ func (txn *Transaction) GetVM(name string) (_ *v1.VM, err error) {
 		return nil, err
 	}
 
+	vm.Version = item.Version()
+
 	return &vm, nil
 }
 
@@ -94,6 +97,8 @@ func (txn *Transaction) ListVMs() (_ []v1.VM, err error) {
 			return nil, err
 		}
 
+		vm.Version = item.Version()
+
 		result = append(result, vm)
 	}
 

internal/controller/store/badger/badger_watch.go

@@ -0,0 +1,175 @@
+package badger
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"errors"
+	"time"
+
+	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
+	"github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/dgraph-io/badger/v3"
+	"github.com/dgraph-io/badger/v3/pb"
+)
+
+func (store *Store) WatchVM(ctx context.Context, vmName string) (chan storepkg.WatchMessage[v1.VM], chan error, error) {
+	readyCh := make(chan struct{}, 1)
+	watchCh := make(chan storepkg.WatchMessage[v1.VM], 1)
+	errCh := make(chan error, 1)
+	subCtx, subCtxCancel := context.WithCancel(ctx)
+
+	go func() {
+		defer subCtxCancel()
+		defer close(watchCh)
+		defer close(errCh)
+
+		var initialVM *v1.VM
+		var checkedInitialVM bool
+
+		if err := store.db.Subscribe(subCtx, func(kvList *badger.KVList) error {
+			if !checkedInitialVM {
+				// Notify the caller that we've subscribed, but don't block,
+				// because we may observe multiple watch barriers, yet
+				// we only need a single barrier to make things work
+				select {
+				case readyCh <- struct{}{}:
+				default:
+				}
+
+				// Now that the subscription has started,
+				// retrieve the initial VM, if any
+				err := store.View(func(txn storepkg.Transaction) error {
+					var err error
+
+					initialVM, err = txn.GetVM(vmName)
+
+					return err
+				})
+				if err != nil && !errors.Is(err, storepkg.ErrNotFound) {
+					return err
+				}
+				if initialVM != nil {
+					notification := storepkg.WatchMessage[v1.VM]{
+						Type:   storepkg.WatchMessageTypeAdded,
+						Object: *initialVM,
+					}
+
+					select {
+					case watchCh <- notification:
+					case <-subCtx.Done():
+						return subCtx.Err()
+					}
+				}
+
+				checkedInitialVM = true
+			}
+
+			for _, kv := range kvList.GetKv() {
+				switch {
+				case bytes.Equal(kv.GetKey(), WatchBarrierKey()):
+					// We only need watch barriers so that the Subscribe()'s callback
+					// is called at least once, thus we can simply do nothing here
+				case bytes.Equal(kv.GetKey(), VMKey(vmName)):
+					// Skip all KVs with versions before or equal
+					// to the initial VM's version, if any
+					if initialVM != nil && kv.GetVersion() <= initialVM.Version {
+						continue
+					}
+
+					if kv.GetValue() == nil {
+						if initialVM == nil {
+							// VM is already deleted
+							continue
+						}
+
+						// VM was deleted
+						notification := storepkg.WatchMessage[v1.VM]{
+							Type: storepkg.WatchMessageTypeDeleted,
+						}
+
+						select {
+						case watchCh <- notification:
+						case <-subCtx.Done():
+							return subCtx.Err()
+						}
+
+						initialVM = nil
+					} else {
+						// VM was created or modified
+						var vm v1.VM
+
+						if err := json.Unmarshal(kv.GetValue(), &vm); err != nil {
+							return err
+						}
+
+						vm.Version = kv.GetVersion()
+
+						var watchMessageType storepkg.WatchMessageType
+
+						if initialVM == nil {
+							watchMessageType = storepkg.WatchMessageTypeAdded
+
+							initialVM = &vm
+						} else {
+							watchMessageType = storepkg.WatchMessageTypeModified
+						}
+
+						notification := storepkg.WatchMessage[v1.VM]{
+							Type:   watchMessageType,
+							Object: vm,
+						}
+
+						select {
+						case watchCh <- notification:
+						case <-subCtx.Done():
+							return subCtx.Err()
+						}
+					}
+				default:
+					// Ignore unexpected keys
+					continue
+				}
+			}
+
+			return nil
+		}, []pb.Match{
+			{
+				Prefix: WatchBarrierKey(),
+			},
+			{
+				Prefix: VMKey(vmName),
+			},
+		}); err != nil {
+			errCh <- err
+		}
+	}()
+
+	// Trigger the watch barrier so that Subscribe() callback gets invoked
+	if err := store.notifyWatchBarrier(); err != nil {
+		subCtxCancel()
+
+		return nil, nil, err
+	}
+
+	// Wait for the Subscribe() callback to be invoked
+Outer:
+	for {
+		select {
+		case <-readyCh:
+			// Subscription has started
+			break Outer
+		case <-time.After(time.Second):
+			// Possible race with late goroutine start, re-issue watch barrier
+			if err := store.notifyWatchBarrier(); err != nil {
+				subCtxCancel()
+
+				return nil, nil, err
+			}
+		case <-ctx.Done():
+			return nil, nil, ctx.Err()
+		}
+	}
+
+	return watchCh, errCh, nil
+}

internal/controller/store/badger/badger_watch_barrier.go

@@ -0,0 +1,18 @@
+package badger
+
+import (
+	"github.com/dgraph-io/badger/v3"
+	"github.com/google/uuid"
+)
+
+const SpaceWatchBarrier = "/watch-barrier"
+
+func WatchBarrierKey() []byte {
+	return []byte(SpaceWatchBarrier)
+}
+
+func (store *Store) notifyWatchBarrier() error {
+	return store.db.Update(func(txn *badger.Txn) error {
+		return txn.Set(WatchBarrierKey(), []byte(uuid.NewString()))
+	})
+}