From 309831c42534ec6742e55cf39c2a7b204b977115 Mon Sep 17 00:00:00 2001
From: luke-lombardi <33990301+luke-lombardi@users.noreply.github.com>
Date: Tue, 7 May 2024 16:02:54 -0400
Subject: [PATCH] Fix machine info fields (#193)

Co-authored-by: Luke Lombardi <luke@beam.cloud>
---
 internal/repository/base.go           |  4 ++--
 internal/repository/provider_redis.go | 22 ++++++++++++++--------
 internal/repository/worker_redis.go   |  1 -
 3 files changed, 16 insertions(+), 11 deletions(-)

diff --git a/internal/repository/base.go b/internal/repository/base.go
index 2540ae5b9..7022d2762 100644
--- 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
index 17ad9eb56..683c79ebe 100644
--- 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
index 2449a7f4b..2f562cfe6 100644
--- 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
 	}