From 219e20fdbad755656d4331cf01e1ee2b3132a09b Mon Sep 17 00:00:00 2001
From: Amir Raminfar <findamir@gmail.com>
Date: Tue, 21 May 2024 10:57:16 -0700
Subject: [PATCH] chore: fixes race condition in container store (#2965)

---
 internal/docker/container_store.go | 27 +++++++++++++++++++--------
 1 file changed, 19 insertions(+), 8 deletions(-)

diff --git a/internal/docker/container_store.go b/internal/docker/container_store.go
index 98c2199d..245d1f84 100644
--- a/internal/docker/container_store.go
+++ b/internal/docker/container_store.go
@@ -128,19 +128,30 @@ func (s *ContainerStore) init() {
 				s.containers.Delete(event.ActorID)
 
 			case "die":
-				if container, ok := s.containers.Load(event.ActorID); ok {
-					log.Debugf("container %s died", container.ID)
-					container.State = "exited"
-				}
+				s.containers.Compute(event.ActorID, func(c *Container, loaded bool) (*Container, bool) {
+					if loaded {
+						log.Debugf("container %s died", c.ID)
+						c.State = "exited"
+						return c, false
+					} else {
+						return c, true
+					}
+				})
 			case "health_status: healthy", "health_status: unhealthy":
 				healthy := "unhealthy"
 				if event.Name == "health_status: healthy" {
 					healthy = "healthy"
 				}
-				if container, ok := s.containers.Load(event.ActorID); ok {
-					log.Debugf("container %s is %s", container.ID, healthy)
-					container.Health = healthy
-				}
+
+				s.containers.Compute(event.ActorID, func(c *Container, loaded bool) (*Container, bool) {
+					if loaded {
+						log.Debugf("health status for container %s is %s", c.ID, healthy)
+						c.Health = healthy
+						return c, false
+					} else {
+						return c, true
+					}
+				})
 			}
 			s.subscribers.Range(func(c context.Context, events chan ContainerEvent) bool {
 				select {