allocate credentials before starting firecracker VM (#320)

synadia-io · Jul 11, 2024 · 7ed512e · 7ed512e
1 parent e114360
commit 7ed512e
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 15 deletions.
diff --git a/agent/agent.go b/agent/agent.go
@@ -26,7 +26,7 @@ import (
 )
 
 const (
-	defaultAgentHandshakeTimeoutMillis  = 500
+	defaultAgentHandshakeTimeoutMillis  = 1000
 	runloopSleepInterval                = 250 * time.Millisecond
 	runloopTickInterval                 = 2500 * time.Millisecond
 	workloadExecutionSleepTimeoutMillis = 1000
@@ -385,14 +385,22 @@ func (a *Agent) initNATS() error {
 	}
 
 	pk, _ := pair.PublicKey()
-	a.nc, err = nats.Connect(url, nats.Nkey(pk, func(b []byte) ([]byte, error) {
-		fmt.Fprintf(os.Stdout, "Attempting to sign NATS server nonce for internal NATS connection; public key: %s", pk)
-		return pair.Sign(b)
-	}))
+	for attempt := 0; attempt < 3; attempt++ {
+		a.nc, err = nats.Connect(url, nats.Nkey(pk, func(b []byte) ([]byte, error) {
+			fmt.Fprintf(os.Stdout, "Attempting to authenticate to internal NATS as %s", pk)
+			return pair.Sign(b)
+		}))
+		if err != nil {
+			time.Sleep(100 * time.Millisecond)
+		} else {
+			break
+		}
+	}
 	if err != nil {
-		fmt.Fprintf(os.Stderr, "failed to connect to shared NATS: %s", err)
+		fmt.Fprintf(os.Stderr, "failed to connect to internal NATS: %s", err)
 		return err
 	}
+
 	fmt.Printf("Connected to internal NATS: %s\n", url)
 
 	js, err := a.nc.JetStream()

diff --git a/internal/node/processmanager/firecracker_procman.go b/internal/node/processmanager/firecracker_procman.go
@@ -14,6 +14,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/rs/xid"
 	agentapi "github.com/synadia-io/nex/internal/agent-api"
 	"github.com/synadia-io/nex/internal/models"
 	internalnats "github.com/synadia-io/nex/internal/node/internal-nats"
@@ -155,19 +156,19 @@ func (f *FirecrackerProcessManager) Start(delegate ProcessDelegate) error {
 				time.Sleep(runloopSleepInterval)
 				continue
 			}
-
-			vm, err := createAndStartVM(context.TODO(), f.config, f.log)
+			vmmID := xid.New().String()
+			workloadKey, err := f.intNats.CreateCredentials(vmmID)
 			if err != nil {
-				f.log.Warn("Failed to create VMM for warming pool.", slog.Any("err", err))
+				f.log.Error("Failed to create workload user", slog.Any("err", err))
 				continue
 			}
+			workloadSeed, _ := workloadKey.Seed()
 
-			workloadKey, err := f.intNats.CreateCredentials(vm.vmmID)
+			vm, err := createAndStartVM(context.TODO(), vmmID, f.config, f.log)
 			if err != nil {
-				f.log.Error("Failed to create workload user", slog.Any("err", err))
+				f.log.Warn("Failed to create VMM for warming pool.", slog.Any("err", err))
 				continue
 			}
-			workloadSeed, _ := workloadKey.Seed()
 
 			err = f.setMetadata(vm, string(workloadSeed))
 			if err != nil {

diff --git a/internal/node/processmanager/running_vm.go b/internal/node/processmanager/running_vm.go
@@ -19,7 +19,6 @@ import (
 
 	"github.com/firecracker-microvm/firecracker-go-sdk"
 	"github.com/firecracker-microvm/firecracker-go-sdk/client/models"
-	"github.com/rs/xid"
 
 	agentapi "github.com/synadia-io/nex/internal/agent-api"
 	nexmodels "github.com/synadia-io/nex/internal/models"
@@ -89,8 +88,7 @@ func (vm *runningFirecracker) shutdown() {
 }
 
 // Create a VMM with a given set of options and start the VM
-func createAndStartVM(ctx context.Context, config *nexmodels.NodeConfiguration, log *slog.Logger) (*runningFirecracker, error) {
-	vmmID := xid.New().String()
+func createAndStartVM(ctx context.Context, vmmID string, config *nexmodels.NodeConfiguration, log *slog.Logger) (*runningFirecracker, error) {
 
 	fcCfg, err := generateFirecrackerConfig(vmmID, config)
 	if err != nil {