Use status files to retrieve PID

dmjb · dmjb · commit 366e4bf39d1e · 2025-09-10T13:00:59.000+01:00
This follows on from the previous PRs which migrated PIDs from the PID
files to status files.
diff --git a/pkg/transport/proxy/manager.go b/pkg/transport/proxy/manager.go
@@ -9,34 +9,6 @@ import (
 // We may want to move these operations behind an interface. For now, they
 // have been moved to this package to keep proxy-related logic grouped together.
 
-// StopProcess stops the proxy process associated with the container
-func StopProcess(containerBaseName string) {
-	if containerBaseName == "" {
-		logger.Warnf("Warning: Could not find base container name in labels")
-		return
-	}
-
-	// Try to read the PID file and kill the process
-	pid, err := process.ReadPIDFile(containerBaseName)
-	if err != nil {
-		logger.Errorf("No PID file found for %s, proxy may not be running in detached mode", containerBaseName)
-		return
-	}
-
-	// PID file found, try to kill the process
-	logger.Infof("Stopping proxy process (PID: %d)...", pid)
-	if err := process.KillProcess(pid); err != nil {
-		logger.Warnf("Warning: Failed to kill proxy process: %v", err)
-	} else {
-		logger.Info("Proxy process stopped")
-	}
-
-	// Remove the PID file
-	if err := process.RemovePIDFile(containerBaseName); err != nil {
-		logger.Warnf("Warning: Failed to remove PID file: %v", err)
-	}
-}
-
 // IsRunning checks if the proxy process is running
 func IsRunning(containerBaseName string) bool {
 	logger.Debugf("Checking if proxy process is running for container %s", containerBaseName)
diff --git a/pkg/workloads/manager.go b/pkg/workloads/manager.go
@@ -570,12 +570,7 @@ func (d *defaultManager) getWorkloadContainer(ctx context.Context, name string)
 func (d *defaultManager) stopProxyIfNeeded(ctx context.Context, name, baseName string) {
 	logger.Infof("Removing proxy process for %s...", name)
 	if baseName != "" {
-		proxy.StopProcess(baseName)
-		// TODO: refactor the StopProcess function to stop dealing explicitly with PID files.
-		// Note that this is not a blocker for k8s since this code path is not called there.
-		if err := d.statuses.ResetWorkloadPID(ctx, baseName); err != nil {
-			logger.Warnf("Warning: Failed to reset workload %s PID: %v", name, err)
-		}
+		d.stopProcess(ctx, baseName)
 	}
 }
 
@@ -962,7 +957,7 @@ func (d *defaultManager) stopSingleContainerWorkload(ctx context.Context, worklo
 
 	name := labels.GetContainerBaseName(workload.Labels)
 	// Stop the proxy process
-	proxy.StopProcess(name)
+	d.stopProcess(ctx, name)
 	// TODO: refactor the StopProcess function to stop dealing explicitly with PID files.
 	// Note that this is not a blocker for k8s since this code path is not called there.
 	if err := d.statuses.ResetWorkloadPID(ctx, name); err != nil {
@@ -1121,3 +1116,36 @@ func (d *defaultManager) getRemoteWorkloadsFromState(
 
 	return remoteWorkloads, nil
 }
+
+// stopProcess stops the proxy process associated with the container
+func (d *defaultManager) stopProcess(ctx context.Context, name string) {
+	if name == "" {
+		logger.Warnf("Warning: Could not find base container name in labels")
+		return
+	}
+
+	// Try to read the PID file and kill the process
+	pid, err := process.ReadPIDFile(name)
+	if err != nil {
+		logger.Errorf("No PID file found for %s, proxy may not be running in detached mode", name)
+		return
+	}
+
+	// PID file found, try to kill the process
+	logger.Infof("Stopping proxy process (PID: %d)...", pid)
+	if err := process.KillProcess(pid); err != nil {
+		logger.Warnf("Warning: Failed to kill proxy process: %v", err)
+	} else {
+		logger.Info("Proxy process stopped")
+	}
+
+	// Remove the PID file
+	if err := process.RemovePIDFile(name); err != nil {
+		logger.Warnf("Warning: Failed to remove PID file: %v", err)
+	}
+	// TODO: refactor the StopProcess function to stop dealing explicitly with PID files.
+	// Note that this is not a blocker for k8s since this code path is not called there.
+	if err := d.statuses.ResetWorkloadPID(ctx, name); err != nil {
+		logger.Warnf("Warning: Failed to reset workload %s PID: %v", name, err)
+	}
+}
diff --git a/pkg/workloads/statuses/file_status.go b/pkg/workloads/statuses/file_status.go
@@ -19,7 +19,6 @@ import (
 	"github.com/stacklok/toolhive/pkg/logger"
 	"github.com/stacklok/toolhive/pkg/process"
 	"github.com/stacklok/toolhive/pkg/state"
-	"github.com/stacklok/toolhive/pkg/transport/proxy"
 	"github.com/stacklok/toolhive/pkg/workloads/types"
 )
 
@@ -114,6 +113,7 @@ type workloadStatusFile struct {
 
 // GetWorkload retrieves the status of a workload by its name.
 func (f *fileStatusManager) GetWorkload(ctx context.Context, workloadName string) (core.Workload, error) {
+	var pid int
 	result := core.Workload{Name: workloadName}
 	fileFound := false
 
@@ -152,6 +152,8 @@ func (f *fileStatusManager) GetWorkload(ctx context.Context, workloadName string
 			}
 		}
 
+		pid = statusFile.ProcessID
+
 		return nil
 	})
 	if err != nil {
@@ -172,7 +174,7 @@ func (f *fileStatusManager) GetWorkload(ctx context.Context, workloadName string
 
 		// If workload is running, validate against runtime
 		if result.Status == rt.WorkloadStatusRunning {
-			return f.validateRunningWorkload(ctx, workloadName, result)
+			return f.validateRunningWorkload(ctx, workloadName, result, pid)
 		}
 
 		// Return file data
@@ -197,7 +199,7 @@ func (f *fileStatusManager) ListWorkloads(ctx context.Context, listAll bool, lab
 	}
 
 	// Get workloads from files
-	fileWorkloads, err := f.getWorkloadsFromFiles()
+	fileWorkloadsWithPID, err := f.getWorkloadsFromFiles()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get workloads from files: %w", err)
 	}
@@ -206,7 +208,7 @@ func (f *fileStatusManager) ListWorkloads(ctx context.Context, listAll bool, lab
 	// There's currently an import cycle between this package and the runconfig package
 
 	// Create a map of runtime workloads by name for easy lookup
-	workloadMap := f.mergeRuntimeAndFileWorkloads(ctx, runtimeContainers, fileWorkloads)
+	workloadMap := f.mergeRuntimeAndFileWorkloads(ctx, runtimeContainers, fileWorkloadsWithPID)
 
 	// Convert map to slice and apply filters
 	var workloads []core.Workload
@@ -553,8 +555,14 @@ func (f *fileStatusManager) getWorkloadFromRuntime(ctx context.Context, workload
 	return types.WorkloadFromContainerInfo(&info)
 }
 
+// workloadWithPID holds a workload and its associated PID for internal processing
+type workloadWithPID struct {
+	workload core.Workload
+	pid      int
+}
+
 // getWorkloadsFromFiles retrieves all workloads from status files.
-func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]core.Workload, error) {
+func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]workloadWithPID, error) {
 	// Ensure base directory exists
 	if err := f.ensureBaseDir(); err != nil {
 		return nil, fmt.Errorf("failed to ensure base directory: %w", err)
@@ -566,7 +574,7 @@ func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]core.Workload, e
 		return nil, fmt.Errorf("failed to list status files: %w", err)
 	}
 
-	workloads := make(map[string]core.Workload)
+	workloads := make(map[string]workloadWithPID)
 	ctx := context.Background() // Create context for file locking
 
 	for _, file := range files {
@@ -614,13 +622,15 @@ func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]core.Workload, e
 			}
 
 			// Check if PID migration is needed
+			pid := statusFile.ProcessID
 			if statusFile.Status == rt.WorkloadStatusRunning && statusFile.ProcessID == 0 {
 				// Try PID migration - the migration function will handle cases
 				// where container info is not available gracefully
 				if migratedPID, wasMigrated := f.migratePIDFromFile(workloadName, nil); wasMigrated {
 					// Update the status file with the migrated PID
 					statusFile.ProcessID = migratedPID
 					statusFile.UpdatedAt = time.Now()
+					pid = migratedPID
 					if err := f.writeStatusFile(statusFilePath, *statusFile); err != nil {
 						logger.Warnf("failed to write migrated PID for workload %s: %v", workloadName, err)
 					} else {
@@ -629,7 +639,10 @@ func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]core.Workload, e
 				}
 			}
 
-			workloads[workloadName] = workload
+			workloads[workloadName] = workloadWithPID{
+				workload: workload,
+				pid:      pid,
+			}
 			return nil
 		})
 
@@ -647,7 +660,7 @@ func (f *fileStatusManager) getWorkloadsFromFiles() (map[string]core.Workload, e
 // validateRunningWorkload validates that a workload marked as running in the file
 // is actually running in the runtime and has a healthy proxy process if applicable.
 func (f *fileStatusManager) validateRunningWorkload(
-	ctx context.Context, workloadName string, result core.Workload,
+	ctx context.Context, workloadName string, result core.Workload, pid int,
 ) (core.Workload, error) {
 	// For remote workloads, we don't need to validate against the container runtime
 	// since they don't have containers
@@ -667,7 +680,7 @@ func (f *fileStatusManager) validateRunningWorkload(
 	}
 
 	// Check if proxy process is running when workload is running
-	if unhealthyWorkload, isUnhealthy := f.checkProxyHealth(ctx, workloadName, result, containerInfo); isUnhealthy {
+	if unhealthyWorkload, isUnhealthy := f.isProxyUnhealthy(ctx, workloadName, result, containerInfo, pid); isUnhealthy {
 		return unhealthyWorkload, nil
 	}
 
@@ -722,19 +735,21 @@ func (f *fileStatusManager) handleRuntimeMissing(
 	return fileWorkload, nil
 }
 
-// checkProxyHealth checks if the proxy process is running for the workload.
+// isProxyUnhealthy checks if the proxy process is running for the workload.
 // Returns (unhealthyWorkload, true) if proxy is not running, (emptyWorkload, false) if proxy is healthy or not applicable.
-func (f *fileStatusManager) checkProxyHealth(
-	ctx context.Context, workloadName string, result core.Workload, containerInfo rt.ContainerInfo,
+func (f *fileStatusManager) isProxyUnhealthy(
+	ctx context.Context, workloadName string, result core.Workload, containerInfo rt.ContainerInfo, pid int,
 ) (core.Workload, bool) {
 	// Use original container labels (before filtering) to get base name
 	baseName := labels.GetContainerBaseName(containerInfo.Labels)
 	if baseName == "" {
 		return core.Workload{}, false // No proxy check needed
 	}
 
-	proxyRunning := proxy.IsRunning(baseName)
-	if proxyRunning {
+	proxyRunning, err := process.FindProcess(pid)
+	if err != nil {
+		logger.Warnf("unable to find process %d: %v", pid, err)
+	} else if proxyRunning {
 		return core.Workload{}, false // Proxy is healthy
 	}
 
@@ -775,7 +790,7 @@ func (*fileStatusManager) mergeHealthyWorkloadData(containerInfo rt.ContainerInf
 // validateWorkloadInList validates a workload during list operations, similar to validateRunningWorkload
 // but with different error handling to avoid disrupting the entire list operation.
 func (f *fileStatusManager) validateWorkloadInList(
-	ctx context.Context, workloadName string, fileWorkload core.Workload, containerInfo rt.ContainerInfo,
+	ctx context.Context, workloadName string, fileWorkload core.Workload, containerInfo rt.ContainerInfo, pid int,
 ) (core.Workload, error) {
 	// Only validate if file shows running status
 	if fileWorkload.Status != rt.WorkloadStatusRunning {
@@ -797,7 +812,7 @@ func (f *fileStatusManager) validateWorkloadInList(
 	}
 
 	// Check if proxy process is running when workload is running
-	if unhealthyWorkload, isUnhealthy := f.checkProxyHealth(ctx, workloadName, fileWorkload, containerInfo); isUnhealthy {
+	if unhealthyWorkload, isUnhealthy := f.isProxyUnhealthy(ctx, workloadName, fileWorkload, containerInfo, pid); isUnhealthy {
 		return unhealthyWorkload, nil
 	}
 
@@ -809,7 +824,7 @@ func (f *fileStatusManager) validateWorkloadInList(
 func (f *fileStatusManager) mergeRuntimeAndFileWorkloads(
 	ctx context.Context,
 	runtimeContainers []rt.ContainerInfo,
-	fileWorkloads map[string]core.Workload,
+	fileWorkloadsWithPID map[string]workloadWithPID,
 ) map[string]core.Workload {
 	runtimeWorkloadMap := make(map[string]rt.ContainerInfo)
 	for _, container := range runtimeContainers {
@@ -840,15 +855,17 @@ func (f *fileStatusManager) mergeRuntimeAndFileWorkloads(
 	}
 
 	// Then, merge with file workloads, validating running workloads
-	for name, fileWorkload := range fileWorkloads {
+	for name, fileWorkloadWithPID := range fileWorkloadsWithPID {
+		fileWorkload := fileWorkloadWithPID.workload
+		pid := fileWorkloadWithPID.pid
 
 		if fileWorkload.Remote { // Remote workloads are not managed by the container runtime
 			workloadMap[name] = fileWorkload
 			continue
 		}
 		if runtimeContainer, exists := runtimeWorkloadMap[name]; exists {
 			// Validate running workloads similar to GetWorkload
-			validatedWorkload, err := f.validateWorkloadInList(ctx, name, fileWorkload, runtimeContainer)
+			validatedWorkload, err := f.validateWorkloadInList(ctx, name, fileWorkload, runtimeContainer, pid)
 			if err != nil {
 				logger.Warnf("failed to validate workload %s in list: %v", name, err)
 				// Fall back to basic merge without validation
@@ -870,8 +887,9 @@ func (f *fileStatusManager) mergeRuntimeAndFileWorkloads(
 			if err != nil {
 				logger.Warnf("failed to handle missing runtime for workload %s: %v", name, err)
 				workloadMap[name] = fileWorkload
+			} else {
+				workloadMap[name] = updatedWorkload
 			}
-			workloadMap[name] = updatedWorkload
 		}
 	}
 	return workloadMap
diff --git a/pkg/workloads/statuses/file_status_test.go b/pkg/workloads/statuses/file_status_test.go
@@ -526,8 +526,8 @@ func TestFileStatusManager_ValidateRunningWorkload_Remote(t *testing.T) {
 	// Mock runtime should NOT be called for remote workloads
 	// (no expectations set, so any call would fail the test)
 
-	// Validate the remote workload
-	result, err := manager.validateRunningWorkload(ctx, "remote-test", remoteWorkload)
+	// Validate the remote workload (PID is irrelevant for remote workloads)
+	result, err := manager.validateRunningWorkload(ctx, "remote-test", remoteWorkload, 0)
 	require.NoError(t, err)
 
 	// Should return the workload unchanged without calling runtime
