public_git/orchard
1
0
Fork 0
mirror of https://github.com/cirruslabs/orchard.git synced 2026-10-19 16:36:52 +02:00
Code Issues Packages Projects Releases Wiki Activity

Refactorings for simplify readability (#35)

Browse Source
This commit is contained in:
Fedor Korotkov
2023-03-17 06:11:28 -04:00
committed by GitHub
parent 3ecf98c039
commit fb3056d3ae
21 changed files with 360 additions and 876 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/controller/api.go
+3 -3
View File
@@ -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_vms_portforward.go
+15 -14
View File
@@ -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/controller.go
+7 -6
View File
@@ -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/notifier/notifier.go
+57
View File
@@ -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/rendezvous/watcher_test.go → internal/controller/notifier/notifier_test.go
+7 -7
View File
@@ -1,9 +1,9 @@
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"
@@ -11,25 +11,25 @@ import (
"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/proxy/proxy.go
+49
View File
@@ -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/rendezvous/proxy_test.go → internal/controller/proxy/proxy_test.go
+3 -3
View File
@@ -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()
internal/controller/rendezvous/concurrentmap/concurrentmap.go
-39
View File
@@ -1,39 +0,0 @@
package concurrentmap
import (
"sync"
)
type ConcurrentMap[T any] struct {
nonConcurrentMap map[any]T
mtx sync.Mutex
}
func NewConcurrentMap[T any]() *ConcurrentMap[T] {
return &ConcurrentMap[T]{
nonConcurrentMap: map[any]T{},
}
}
func (cmap *ConcurrentMap[T]) Load(key any) (T, bool) {
cmap.mtx.Lock()
defer cmap.mtx.Unlock()
result, ok := cmap.nonConcurrentMap[key]
return result, ok
}
func (cmap *ConcurrentMap[T]) Store(id any, value T) {
cmap.mtx.Lock()
defer cmap.mtx.Unlock()
cmap.nonConcurrentMap[id] = value
}
func (cmap *ConcurrentMap[T]) Delete(key any) {
cmap.mtx.Lock()
defer cmap.mtx.Unlock()
delete(cmap.nonConcurrentMap, key)
}
internal/controller/rendezvous/proxy.go
-49
View File
@@ -1,49 +0,0 @@
package rendezvous
import (
"context"
"errors"
"github.com/cirruslabs/orchard/internal/controller/rendezvous/concurrentmap"
"net"
)
var ErrInvalidToken = errors.New("invalid proxy token")
type Proxy struct {
tokens *concurrentmap.ConcurrentMap[*TokenSlot]
}
type TokenSlot struct {
ctx context.Context
ch chan net.Conn
}
func NewProxy() *Proxy {
return &Proxy{
tokens: concurrentmap.NewConcurrentMap[*TokenSlot](),
}
}
func (proxy *Proxy) Request(ctx context.Context, token string) (chan net.Conn, func()) {
tokenSlot := &TokenSlot{
ctx: ctx,
ch: make(chan net.Conn),
}
proxy.tokens.Store(token, tokenSlot)
return tokenSlot.ch, func() {
proxy.tokens.Delete(token)
}
}
func (proxy *Proxy) Respond(token string, conn net.Conn) (context.Context, error) {
tokenSlot, ok := proxy.tokens.Load(token)
if !ok {
return nil, ErrInvalidToken
}
tokenSlot.ch <- conn
return tokenSlot.ctx, nil
}
internal/controller/rendezvous/watcher.go
-57
View File
@@ -1,57 +0,0 @@
package rendezvous
import (
"context"
"errors"
"fmt"
"github.com/cirruslabs/orchard/internal/controller/rendezvous/concurrentmap"
"github.com/cirruslabs/orchard/rpc"
)
var ErrNoWorker = errors.New("no worker registered with this name")
type Watcher struct {
workers *concurrentmap.ConcurrentMap[*WorkerSlot]
}
type WorkerSlot struct {
ctx context.Context
ch chan *rpc.WatchFromController
}
func NewWatcher() *Watcher {
return &Watcher{
workers: concurrentmap.NewConcurrentMap[*WorkerSlot](),
}
}
func (watcher *Watcher) Subscribe(ctx context.Context, worker string) (chan *rpc.WatchFromController, func()) {
subCtx, cancel := context.WithCancel(ctx)
workerCh := make(chan *rpc.WatchFromController)
watcher.workers.Store(worker, &WorkerSlot{
ctx: subCtx,
ch: workerCh,
})
return workerCh, func() {
watcher.workers.Delete(worker)
cancel()
}
}
func (watcher *Watcher) Notify(ctx context.Context, topic string, msg *rpc.WatchFromController) error {
slot, ok := watcher.workers.Load(topic)
if !ok {
return fmt.Errorf("%w: %s", ErrNoWorker, topic)
}
select {
case slot.ch <- msg:
return nil
case <-slot.ctx.Done():
return slot.ctx.Err()
case <-ctx.Done():
return ctx.Err()
}
}
internal/controller/rpc.go
+19 -36
View File
@@ -3,6 +3,9 @@ package controller
import (
v1pkg "github.com/cirruslabs/orchard/pkg/resource/v1"
"github.com/cirruslabs/orchard/rpc"
"google.golang.org/grpc/metadata"
"google.golang.org/protobuf/types/known/emptypb"
//nolint:staticcheck // https://github.com/mitchellh/go-grpc-net-conn/pull/1
"github.com/golang/protobuf/proto"
grpc_net_conn "github.com/mitchellh/go-grpc-net-conn"
@@ -10,24 +13,18 @@ import (
"google.golang.org/grpc/status"
)
func (controller *Controller) Watch(stream rpc.Controller_WatchServer) error {
func (controller *Controller) Watch(_ *emptypb.Empty, stream rpc.Controller_WatchServer) error {
if !controller.authorizeGRPC(stream.Context(), v1pkg.ServiceAccountRoleWorker) {
return status.Errorf(codes.Unauthenticated, "auth failed")
}
// The first message is always an initialization
watchFromWorker, err := stream.Recv()
if err != nil {
return err
workerUIDMetadataValue := metadata.ValueFromIncomingContext(stream.Context(), rpc.MetadataWorkerUIDKey)
if len(workerUIDMetadataValue) == 0 {
return status.Errorf(codes.InvalidArgument, "no worker UID in metadata")
}
initAction, ok := watchFromWorker.Action.(*rpc.WatchFromWorker_InitAction)
if !ok {
return status.Errorf(codes.FailedPrecondition, "expected an initialization message")
}
// Subscribe to rendez-vous requests from the API for this worker
workerCh, cancel := controller.watcher.Subscribe(stream.Context(), initAction.InitAction.WorkerUid)
workerUID := workerUIDMetadataValue[0]
workerCh, cancel := controller.workerNotifier.Register(stream.Context(), workerUID)
defer cancel()
for {
@@ -47,39 +44,25 @@ func (controller *Controller) PortForward(stream rpc.Controller_PortForwardServe
return status.Errorf(codes.Unauthenticated, "auth failed")
}
// The first message is always an initialization
portForwardFromWorker, err := stream.Recv()
if err != nil {
return err
sessionMetadataValue := metadata.ValueFromIncomingContext(stream.Context(), rpc.MetadataWorkerPortForwardingSessionKey)
if len(sessionMetadataValue) == 0 {
return status.Errorf(codes.InvalidArgument, "no session in metadata")
}
initAction, ok := portForwardFromWorker.Action.(*rpc.PortForwardFromWorker_InitAction)
if !ok {
return status.Errorf(codes.FailedPrecondition, "expected an initialization message")
}
// Satisfy the rendez-vous request
conn := &grpc_net_conn.Conn{
Stream: stream,
Request: &rpc.PortForwardFromController{
Action: &rpc.PortForwardFromController_DataAction{
DataAction: &rpc.PortForwardFromController_Data{},
},
},
Response: &rpc.PortForwardFromWorker{
Action: &rpc.PortForwardFromWorker_DataAction{
DataAction: &rpc.PortForwardFromWorker_Data{},
},
},
Stream: stream,
Request: &rpc.PortForwardData{},
Response: &rpc.PortForwardData{},
Encode: grpc_net_conn.SimpleEncoder(func(message proto.Message) *[]byte {
return &message.(*rpc.PortForwardFromController).Action.(*rpc.PortForwardFromController_DataAction).DataAction.Data
return &message.(*rpc.PortForwardData).Data
}),
Decode: grpc_net_conn.SimpleDecoder(func(message proto.Message) *[]byte {
return &message.(*rpc.PortForwardFromWorker).Action.(*rpc.PortForwardFromWorker_DataAction).DataAction.Data
return &message.(*rpc.PortForwardData).Data
}),
}
proxyCtx, err := controller.proxy.Respond(initAction.InitAction.Token, conn)
// make proxy aware of the connection
proxyCtx, err := controller.proxy.Respond(sessionMetadataValue[0], conn)
if err != nil {
return err
}
internal/controller/scheduler.go
+2 -2
View File
@@ -51,7 +51,7 @@ func runSchedulerInner(store storepkg.Store) error {
for _, vm := range vms {
vm := vm
if vm.Worker != "" {
if vm.WorkerUID != "" {
continue
}
@@ -59,7 +59,7 @@ func runSchedulerInner(store storepkg.Store) error {
for _, worker := range workers {
worker := worker
vm.Worker = worker.Name
vm.WorkerUID = worker.UID
err := store.Update(func(txn storepkg.Transaction) error {
return txn.SetVM(vm)