Refactorings for simplify readability (#35)

Author: fkorotkov

internal/controller/rendezvous/concurrentmap/concurrentmap.go renamed to internal/concurrentmap/concurrentmap.go

@@ -6,8 +6,8 @@ import (
 	"errors"
 	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
 	"github.com/cirruslabs/orchard/internal/responder"
-	"github.com/cirruslabs/orchard/pkg/client"
 	v1pkg "github.com/cirruslabs/orchard/pkg/resource/v1"
+	"github.com/cirruslabs/orchard/rpc"
 	"github.com/deckarep/golang-set/v2"
 	"github.com/gin-gonic/gin"
 	"google.golang.org/grpc/metadata"
@@ -169,11 +169,11 @@ func (controller *Controller) authorizeGRPC(ctx context.Context, scopes ...v1pkg
 		return true
 	}
 
-	name := metadata.ValueFromIncomingContext(ctx, client.MetadataServiceAccountNameKey)
+	name := metadata.ValueFromIncomingContext(ctx, rpc.MetadataServiceAccountNameKey)
 	if len(name) != 1 {
 		return false
 	}
-	token := metadata.ValueFromIncomingContext(ctx, client.MetadataServiceAccountTokenKey)
+	token := metadata.ValueFromIncomingContext(ctx, rpc.MetadataServiceAccountTokenKey)
 	if len(token) != 1 {
 		return false
 	}

internal/controller/api.go

@@ -45,35 +45,36 @@ func (controller *Controller) portForwardVM(ctx *gin.Context) responder.Responde
 	}
 
 	// Sanity-check
-	if vm.Worker == "" {
+	if vm.WorkerUID == "" {
 		return responder.Code(http.StatusServiceUnavailable)
 	}
 
-	// Request and wait for a rendez-vous with a worker
-	token := uuid.New().String()
-
-	rendezvousConnCh, cancel := controller.proxy.Request(ctx, token)
+	// Request and wait for a connection with a worker
+	session := uuid.New().String()
+	boomerangConnCh, cancel := controller.proxy.Request(ctx, session)
 	defer cancel()
 
-	err = controller.watcher.Notify(ctx, vm.Worker, &rpc.WatchFromController{
-		Action: &rpc.WatchFromController_PortForwardAction{
-			PortForwardAction: &rpc.WatchFromController_PortForward{
-				Token:  token,
-				VmUid:  vm.UID,
-				VmPort: uint32(port),
+	// send request to worker to initiate port-forwarding connection back to us
+	err = controller.workerNotifier.Notify(ctx, vm.WorkerUID, &rpc.WatchInstruction{
+		Action: &rpc.WatchInstruction_PortForwardAction{
+			PortForwardAction: &rpc.WatchInstruction_PortForward{
+				Session: session,
+				VmUid:   vm.UID,
+				VmPort:  uint32(port),
 			},
 		},
 	})
 	if err != nil {
-		controller.logger.Warnf("failed to rendez-vous with the worker %s: %v", vm.Worker, err)
+		controller.logger.Warnf("failed to request port-forwarding from the worker %s: %v", vm.WorkerUID, err)
 
 		return responder.Code(http.StatusServiceUnavailable)
 	}
 
+	// worker will asynchronously start port-forwarding so we wait
 	select {
-	case rendezvousConn := <-rendezvousConnCh:
+	case fromWorkerConnection := <-boomerangConnCh:
 		websocket.Handler(func(wsConn *websocket.Conn) {
-			if err := proxy.Connections(wsConn, rendezvousConn); err != nil {
+			if err := proxy.Connections(wsConn, fromWorkerConnection); err != nil {
 				controller.logger.Warnf("failed to port-forward: %v", err)
 			}
 		}).ServeHTTP(ctx.Writer, ctx.Request)

internal/controller/api_vms_portforward.go

@@ -5,7 +5,8 @@ import (
 	"crypto/tls"
 	"errors"
 	"fmt"
-	"github.com/cirruslabs/orchard/internal/controller/rendezvous"
+	"github.com/cirruslabs/orchard/internal/controller/notifier"
+	"github.com/cirruslabs/orchard/internal/controller/proxy"
 	storepkg "github.com/cirruslabs/orchard/internal/controller/store"
 	"github.com/cirruslabs/orchard/internal/controller/store/badger"
 	"github.com/cirruslabs/orchard/internal/netconstants"
@@ -37,16 +38,16 @@ type Controller struct {
 	store                storepkg.Store
 	logger               *zap.SugaredLogger
 	grpcServer           *grpc.Server
-	watcher              *rendezvous.Watcher
-	proxy                *rendezvous.Proxy
+	workerNotifier       *notifier.Notifier
+	proxy                *proxy.Proxy
 
 	rpc.UnimplementedControllerServer
 }
 
 func New(opts ...Option) (*Controller, error) {
 	controller := &Controller{
-		watcher: rendezvous.NewWatcher(),
-		proxy:   rendezvous.NewProxy(),
+		workerNotifier: notifier.NewNotifier(),
+		proxy:          proxy.NewProxy(),
 	}
 
 	// Apply options
@@ -129,7 +130,7 @@ func (controller *Controller) EnsureServiceAccount(serviceAccount *v1.ServiceAcc
 
 func (controller *Controller) Run(ctx context.Context) error {
 	// Run the scheduler so that each VM will eventually
-	// be assigned to a specific Worker
+	// be assigned to a specific WorkerUID
 	go func() {
 		err := runScheduler(controller.store)
 		if err != nil {

internal/controller/controller.go

@@ -0,0 +1,57 @@
+package notifier
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"github.com/cirruslabs/orchard/internal/concurrentmap"
+	"github.com/cirruslabs/orchard/rpc"
+)
+
+var ErrNoWorker = errors.New("no worker registered with this name")
+
+type Notifier struct {
+	workers *concurrentmap.ConcurrentMap[*WorkerSlot]
+}
+
+type WorkerSlot struct {
+	ctx context.Context
+	ch  chan *rpc.WatchInstruction
+}
+
+func NewNotifier() *Notifier {
+	return &Notifier{
+		workers: concurrentmap.NewConcurrentMap[*WorkerSlot](),
+	}
+}
+
+func (watcher *Notifier) Register(ctx context.Context, workerUID string) (chan *rpc.WatchInstruction, func()) {
+	subCtx, cancel := context.WithCancel(ctx)
+	workerCh := make(chan *rpc.WatchInstruction)
+
+	watcher.workers.Store(workerUID, &WorkerSlot{
+		ctx: subCtx,
+		ch:  workerCh,
+	})
+
+	return workerCh, func() {
+		watcher.workers.Delete(workerUID)
+		cancel()
+	}
+}
+
+func (watcher *Notifier) Notify(ctx context.Context, workerUID string, msg *rpc.WatchInstruction) error {
+	slot, ok := watcher.workers.Load(workerUID)
+	if !ok {
+		return fmt.Errorf("%w: %s", ErrNoWorker, workerUID)
+	}
+
+	select {
+	case slot.ch <- msg:
+		return nil
+	case <-slot.ctx.Done():
+		return slot.ctx.Err()
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}

internal/controller/notifier/notifier.go

@@ -1,35 +1,35 @@
-package rendezvous_test
+package notifier_test
 
 import (
 	"context"
 	"fmt"
-	"github.com/cirruslabs/orchard/internal/controller/rendezvous"
+	"github.com/cirruslabs/orchard/internal/controller/notifier"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"sync"
 	"testing"
 	"time"
 )
 
-func TestWatcher(t *testing.T) {
+func TestNotifier(t *testing.T) {
 	ctx := context.Background()
 
-	watcher := rendezvous.NewWatcher()
+	notifier := notifier.NewNotifier()
 
 	var topic = uuid.New().String()
 
-	msgCh, cancel := watcher.Subscribe(context.Background(), topic)
+	msgCh, cancel := notifier.Register(context.Background(), topic)
 	defer cancel()
 
 	var wg sync.WaitGroup
 	wg.Add(1)
 
 	go func() {
-		require.NoError(t, watcher.Notify(ctx, topic, nil))
+		require.NoError(t, notifier.Notify(ctx, topic, nil))
 
 		time.Sleep(time.Second)
 
-		require.NoError(t, watcher.Notify(ctx, topic, nil))
+		require.NoError(t, notifier.Notify(ctx, topic, nil))
 
 		wg.Done()
 	}()

internal/controller/rendezvous/watcher_test.go renamed to internal/controller/notifier/notifier_test.go

@@ -0,0 +1,49 @@
+package proxy
+
+import (
+	"context"
+	"errors"
+	"github.com/cirruslabs/orchard/internal/concurrentmap"
+	"net"
+)
+
+var ErrInvalidToken = errors.New("invalid proxy token")
+
+type Proxy struct {
+	sessions *concurrentmap.ConcurrentMap[*TokenSlot]
+}
+
+type TokenSlot struct {
+	ctx context.Context
+	ch  chan net.Conn
+}
+
+func NewProxy() *Proxy {
+	return &Proxy{
+		sessions: concurrentmap.NewConcurrentMap[*TokenSlot](),
+	}
+}
+
+func (proxy *Proxy) Request(ctx context.Context, session string) (chan net.Conn, func()) {
+	tokenSlot := &TokenSlot{
+		ctx: ctx,
+		ch:  make(chan net.Conn),
+	}
+
+	proxy.sessions.Store(session, tokenSlot)
+
+	return tokenSlot.ch, func() {
+		proxy.sessions.Delete(session)
+	}
+}
+
+func (proxy *Proxy) Respond(session string, conn net.Conn) (context.Context, error) {
+	tokenSlot, ok := proxy.sessions.Load(session)
+	if !ok {
+		return nil, ErrInvalidToken
+	}
+
+	tokenSlot.ch <- conn
+
+	return tokenSlot.ctx, nil
+}

internal/controller/proxy/proxy.go

@@ -1,8 +1,8 @@
-package rendezvous_test
+package proxy_test
 
 import (
 	"context"
-	"github.com/cirruslabs/orchard/internal/controller/rendezvous"
+	"github.com/cirruslabs/orchard/internal/controller/proxy"
 	"github.com/google/uuid"
 	"github.com/stretchr/testify/require"
 	"net"
@@ -15,7 +15,7 @@ func TestProxy(t *testing.T) {
 
 	expectedConn, _ := net.Pipe()
 
-	proxy := rendezvous.NewProxy()
+	proxy := proxy.NewProxy()
 
 	token := uuid.New().String()