From be04efd28c54a0d67cdc3fef4017920a90667681 Mon Sep 17 00:00:00 2001
From: Yondon Fu <yondon.fu@gmail.com>
Date: Thu, 25 Jan 2024 19:28:16 +0000
Subject: [PATCH] worker: Use mutex for container management

---
 worker/worker.go | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/worker/worker.go b/worker/worker.go
index a9feeb890..4183fe51e 100644
--- a/worker/worker.go
+++ b/worker/worker.go
@@ -9,6 +9,7 @@ import (
 	"log/slog"
 	"math"
 	"mime/multipart"
+	"sync"
 	"time"
 
 	"github.com/docker/cli/opts"
@@ -45,6 +46,7 @@ type Worker struct {
 
 	dockerClient *client.Client
 	containers   map[string]*RunnerContainer
+	mu           *sync.Mutex
 }
 
 func NewWorker(containerImageID string, gpus []string, modelDir string) (*Worker, error) {
@@ -60,6 +62,7 @@ func NewWorker(containerImageID string, gpus []string, modelDir string) (*Worker
 		modelDir:         modelDir,
 		dockerClient:     dockerClient,
 		containers:       make(map[string]*RunnerContainer),
+		mu:               &sync.Mutex{},
 	}, nil
 }
 
@@ -200,6 +203,9 @@ func (w *Worker) Stop(ctx context.Context, containerName string) error {
 }
 
 func (w *Worker) getWarmContainer(ctx context.Context, containerName string, modelID string) (*RunnerContainer, error) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+
 	filters := filters.NewArgs(filters.Arg("name", "^"+containerName+"$"), filters.Arg("status", "running"))
 	containers, err := w.dockerClient.ContainerList(ctx, types.ContainerListOptions{Filters: filters})
 	if err != nil {