Fix machine info fields (#193)

Co-authored-by: Luke Lombardi <[email protected]>
beam-cloud · May 7, 2024 · 309831c · 309831c
1 parent 7f13e3b
commit 309831c
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 11 deletions.
diff --git a/internal/repository/base.go b/internal/repository/base.go
@@ -101,9 +101,9 @@ type TaskRepository interface {
 
 type ProviderRepository interface {
 	GetMachine(providerName, poolName, machineId string) (*types.ProviderMachineState, error)
-	AddMachine(providerName, poolName, machineId string, info *types.ProviderMachineState) error
+	AddMachine(providerName, poolName, machineId string, machineInfo *types.ProviderMachineState) error
 	RemoveMachine(providerName, poolName, machineId string) error
-	RegisterMachine(providerName, poolName, machineId string, info *types.ProviderMachineState) error
+	RegisterMachine(providerName, poolName, machineId string, newMachineInfo *types.ProviderMachineState) error
 	WaitForMachineRegistration(providerName, poolName, machineId string) (*types.ProviderMachineState, error)
 	ListAllMachines(providerName, poolName string) ([]*types.ProviderMachine, error)
 	SetMachineLock(providerName, poolName, machineId string) error

diff --git a/internal/repository/provider_redis.go b/internal/repository/provider_redis.go
@@ -140,14 +140,14 @@ func (r *ProviderRedisRepository) WaitForMachineRegistration(providerName, poolN
 	}
 }
 
-func (r *ProviderRedisRepository) AddMachine(providerName, poolName, machineId string, info *types.ProviderMachineState) error {
+func (r *ProviderRedisRepository) AddMachine(providerName, poolName, machineId string, machineInfo *types.ProviderMachineState) error {
 	stateKey := common.RedisKeys.ProviderMachineState(providerName, poolName, machineId)
 
-	info.MachineId = machineId
-	info.Status = types.MachineStatusPending
+	machineInfo.MachineId = machineId
+	machineInfo.Status = types.MachineStatusPending
 
 	err := r.rdb.HSet(context.TODO(),
-		stateKey, common.ToSlice(info)).Err()
+		stateKey, common.ToSlice(machineInfo)).Err()
 
 	if err != nil {
 		return fmt.Errorf("failed to set machine state <%v>: %w", stateKey, err)
@@ -179,13 +179,19 @@ func (r *ProviderRedisRepository) RemoveMachine(providerName, poolName, machineI
 	return nil
 }
 
-func (r *ProviderRedisRepository) RegisterMachine(providerName, poolName, machineId string, info *types.ProviderMachineState) error {
+func (r *ProviderRedisRepository) RegisterMachine(providerName, poolName, machineId string, newMachineInfo *types.ProviderMachineState) error {
 	stateKey := common.RedisKeys.ProviderMachineState(providerName, poolName, machineId)
 
-	info.MachineId = machineId
-	info.Status = types.MachineStatusRegistered
+	machineInfo, err := r.getMachineFromKey(stateKey)
+	if err != nil {
+		return fmt.Errorf("failed to get machine state <%v>: %w", stateKey, err)
+	}
+
+	machineInfo.HostName = newMachineInfo.HostName
+	machineInfo.Token = newMachineInfo.Token
+	machineInfo.Status = types.MachineStatusRegistered
 
-	err := r.rdb.HSet(context.TODO(), stateKey, common.ToSlice(info)).Err()
+	err = r.rdb.HSet(context.TODO(), stateKey, common.ToSlice(machineInfo)).Err()
 	if err != nil {
 		return fmt.Errorf("failed to set machine state <%v>: %w", stateKey, err)
 	}

diff --git a/internal/repository/worker_redis.go b/internal/repository/worker_redis.go
@@ -329,7 +329,6 @@ func (r *WorkerRedisRepository) ScheduleContainerRequest(worker *types.Worker, r
 	// Update the worker capacity first
 	err = r.UpdateWorkerCapacity(worker, request, types.RemoveCapacity)
 	if err != nil {
-		log.Println(err)
 		return err
 	}