Adding context instead of if statement

Signed-off-by: Amit Schendel <[email protected]>
kubescape · Jun 5, 2024 · 3db3807 · 3db3807
1 parent 7eb0aae
commit 3db3807
Showing 1 changed file with 20 additions and 16 deletions.
diff --git a/pkg/containerwatcher/v1/container_watcher_private.go b/pkg/containerwatcher/v1/container_watcher_private.go
@@ -31,6 +31,8 @@ func (ch *IGContainerWatcher) containerCallback(notif containercollection.PubSub
 		ch.timeBasedContainers.Add(notif.Container.Runtime.ContainerID)
 	}
 
+	tracingContext := context.Background()
+
 	switch notif.Type {
 	case containercollection.EventTypeAddContainer:
 		logger.L().Info("start monitor on container", helpers.String("container ID", notif.Container.Runtime.ContainerID), helpers.String("k8s workload", k8sContainerID))
@@ -41,28 +43,29 @@ func (ch *IGContainerWatcher) containerCallback(notif containercollection.PubSub
 				logger.L().Fatal("attaching container to syscall tracer", helpers.String("container ID", notif.Container.Runtime.ContainerID), helpers.String("k8s workload", k8sContainerID), helpers.Error(err))
 			}
 
-			// Read the syscall tracer events in a separate goroutine.
+			// Read the syscall tracer events in a separate goroutine until the container is removed (signaled from the tracing context).
 			go func() {
 				for {
-					if !ch.timeBasedContainers.Contains(notif.Container.Runtime.ContainerID) && !ch.preRunningContainersIDs.Contains(notif.Container.Runtime.ContainerID) {
-						logger.L().Info("stop monitor on container - container has been removed",
+					select {
+					case <-tracingContext.Done():
+						logger.L().Info("stop monitor on container - container has terminated",
 							helpers.String("container ID", notif.Container.Runtime.ContainerID),
 							helpers.String("k8s workload", k8sContainerID))
 						return
+					default:
+						evs, err := ch.syscallTracer.Read(notif.Container.Runtime.ContainerID)
+						if err != nil {
+							logger.L().Debug("syscalls perf buffer closed", helpers.String("error", err.Error()))
+							return
+						}
+						for _, ev := range evs {
+							ev.SetContainerMetadata(notif.Container)
+							ch.syscallEventCallback(ev)
+						}
+
+						// Sleep for a while before reading the next batch of events.
+						time.Sleep(2 * time.Second) // TODO: make this configurable.
 					}
-
-					evs, err := ch.syscallTracer.Read(notif.Container.Runtime.ContainerID)
-					if err != nil {
-						logger.L().Debug("syscalls perf buffer closed", helpers.String("error", err.Error()))
-						return
-					}
-					for _, ev := range evs {
-						ev.SetContainerMetadata(notif.Container)
-						ch.syscallEventCallback(ev)
-					}
-
-					// Sleep for a while before reading the next batch of events.
-					time.Sleep(2 * time.Second) // TODO: make this configurable.
 				}
 			}()
 		}
@@ -80,6 +83,7 @@ func (ch *IGContainerWatcher) containerCallback(notif containercollection.PubSub
 		logger.L().Info("stop monitor on container - container has terminated",
 			helpers.String("container ID", notif.Container.Runtime.ContainerID),
 			helpers.String("k8s workload", k8sContainerID))
+		tracingContext.Done()
 		ch.preRunningContainersIDs.Remove(notif.Container.Runtime.ContainerID)
 		ch.timeBasedContainers.Remove(notif.Container.Runtime.ContainerID)
 		ch.syscallTracer.Detach(notif.Container.Mntns)