fix store handling of many pids per container, add tests

Author: grcevski

pkg/components/kube/store.go

@@ -82,15 +82,15 @@ type Store struct {
 
 	metadataNotifier meta.Notifier
 
-	containerIDs map[string]*container.Info
+	containerIDs maps.Map2[string, uint32, *container.Info]
 
 	// stores container info by PID. It is only required for
 	// deleting entries in namespaces and podsByContainer when DeleteProcess is called
 	containerByPID map[uint32]*container.Info
 
 	// a single namespace will point to any container inside the pod
 	// but we don't care which one
-	namespaces map[uint32]*container.Info
+	namespaces maps.Map2[uint32, uint32, *container.Info]
 
 	// container ID to pod matcher
 	podsByContainer map[string]*CachedObjMeta
@@ -126,8 +126,8 @@ func NewStore(kubeMetadata meta.Notifier, resourceLabels ResourceLabels, service
 
 	db := &Store{
 		log:                 log,
-		containerIDs:        map[string]*container.Info{},
-		namespaces:          map[uint32]*container.Info{},
+		containerIDs:        maps.Map2[string, uint32, *container.Info]{},
+		namespaces:          maps.Map2[uint32, uint32, *container.Info]{},
 		podsByContainer:     map[string]*CachedObjMeta{},
 		containerByPID:      map[uint32]*container.Info{},
 		objectMetaByIP:      map[string]*CachedObjMeta{},
@@ -197,13 +197,19 @@ func (s *Store) cacheResourceMetadata(meta *informer.ObjectMeta) *CachedObjMeta
 func (s *Store) On(event *informer.Event) error {
 	switch event.Type {
 	case informer.EventType_CREATED:
+		// go func() {
+		// 	time.Sleep(10 * time.Second)
 		s.addObjectMeta(event.Resource)
+		s.Notify(event)
+		// }()
 	case informer.EventType_UPDATED:
 		s.updateObjectMeta(event.Resource)
+		s.Notify(event)
 	case informer.EventType_DELETED:
 		s.deleteObjectMeta(event.Resource)
+		s.Notify(event)
 	}
-	s.Notify(event)
+	//s.Notify(event)
 	return nil
 }
 
@@ -221,8 +227,8 @@ func (s *Store) AddProcess(pid uint32) {
 
 	s.access.Lock()
 	defer s.access.Unlock()
-	s.namespaces[ifp.PIDNamespace] = &ifp
-	s.containerIDs[ifp.ContainerID] = &ifp
+	s.namespaces.Put(ifp.PIDNamespace, pid, &ifp)
+	s.containerIDs.Put(ifp.ContainerID, pid, &ifp)
 	s.containerByPID[pid] = &ifp
 }
 
@@ -234,8 +240,8 @@ func (s *Store) DeleteProcess(pid uint32) {
 		return
 	}
 	delete(s.containerByPID, pid)
-	delete(s.namespaces, info.PIDNamespace)
-	delete(s.containerIDs, info.ContainerID)
+	s.namespaces.Delete(info.PIDNamespace, pid)
+	s.containerIDs.Delete(info.ContainerID, pid)
 }
 
 func (s *Store) addObjectMeta(meta *informer.ObjectMeta) {
@@ -279,9 +285,11 @@ func (s *Store) unlockedAddObjectMeta(meta *informer.ObjectMeta) {
 		for _, c := range meta.Pod.Containers {
 			s.podsByContainer[c.Id] = cmeta
 			// TODO: make sure we can handle when the containerIDs is set after this function is triggered
-			info, ok := s.containerIDs[c.Id]
+			infos, ok := s.containerIDs[c.Id]
 			if ok {
-				s.namespaces[info.PIDNamespace] = info
+				for pid, info := range infos {
+					s.namespaces.Put(info.PIDNamespace, pid, info)
+				}
 			}
 			s.containersByOwner.Put(oID, c.Id, c)
 		}
@@ -315,10 +323,14 @@ func (s *Store) unlockedDeleteObjectMeta(meta *informer.ObjectMeta) {
 		s.log.Debug("deleting pod from store",
 			"ips", meta.Ips, "pod", meta.Name, "namespace", meta.Namespace, "containers", meta.Pod.Containers)
 		for _, c := range meta.Pod.Containers {
-			info, ok := s.containerIDs[c.Id]
+			infos, ok := s.containerIDs[c.Id]
 			if ok {
-				delete(s.containerIDs, c.Id)
-				delete(s.namespaces, info.PIDNamespace)
+				s.containerIDs.DeleteAll(c.Id)
+				for _, info := range infos {
+					s.namespaces.DeleteAll(info.PIDNamespace)
+					// delete all needs to be done only once, we could alternatively delete one by one pid
+					break
+				}
 			}
 			delete(s.podsByContainer, c.Id)
 			s.containersByOwner.Delete(oID, c.Id)
@@ -345,14 +357,19 @@ func (s *Store) PodByContainerID(cid string) *CachedObjMeta {
 func (s *Store) PodContainerByPIDNs(pidns uint32) (*CachedObjMeta, string) {
 	s.access.RLock()
 	defer s.access.RUnlock()
-	if info, ok := s.namespaces[pidns]; ok {
-		if om, ok := s.podsByContainer[info.ContainerID]; ok {
-			oID := fetchOwnerID(om.Meta)
-			containerName := ""
-			if containerInfo, ok := s.containersByOwner.Get(oID, info.ContainerID); ok {
-				containerName = containerInfo.Name
+	if infos, ok := s.namespaces[pidns]; ok {
+		for _, info := range infos {
+			if om, ok := s.podsByContainer[info.ContainerID]; ok {
+				oID := fetchOwnerID(om.Meta)
+				containerName := ""
+				if containerInfo, ok := s.containersByOwner.Get(oID, info.ContainerID); ok {
+					containerName = containerInfo.Name
+				}
+				return om, containerName
 			}
-			return om, containerName
+			// we break here, the namespace is the same for all pids in the container
+			// we need to check one only
+			break
 		}
 	}
 	return nil, ""