From 49c1fd27acb3f44ebc386aba1163c130cec95510 Mon Sep 17 00:00:00 2001
From: Mrunal Patel <mrunalp@gmail.com>
Date: Fri, 1 Sep 2017 08:49:08 -0700
Subject: [PATCH] server: Add infra container store to track them separately

Signed-off-by: Mrunal Patel <mrunalp@gmail.com>
---
 libkpod/container_server.go | 31 +++++++++++++++++++++++++++----
 server/server.go            | 12 ++++++++++++
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/libkpod/container_server.go b/libkpod/container_server.go
index 68d1a8b1..c51ea4b5 100644
--- a/libkpod/container_server.go
+++ b/libkpod/container_server.go
@@ -143,8 +143,9 @@ func New(config *Config) (*ContainerServer, error) {
 		imageContext:         &types.SystemContext{SignaturePolicyPath: config.SignaturePolicyPath},
 		stateLock:            lock,
 		state: &containerServerState{
-			containers: oci.NewMemoryStore(),
-			sandboxes:  make(map[string]*sandbox.Sandbox),
+			containers:      oci.NewMemoryStore(),
+			infraContainers: oci.NewMemoryStore(),
+			sandboxes:       make(map[string]*sandbox.Sandbox),
 		},
 		config: config,
 	}, nil
@@ -579,8 +580,9 @@ func (c *ContainerServer) Shutdown() error {
 }
 
 type containerServerState struct {
-	containers oci.ContainerStorer
-	sandboxes  map[string]*sandbox.Sandbox
+	containers      oci.ContainerStorer
+	infraContainers oci.ContainerStorer
+	sandboxes       map[string]*sandbox.Sandbox
 }
 
 // AddContainer adds a container to the container state store
@@ -592,6 +594,13 @@ func (c *ContainerServer) AddContainer(ctr *oci.Container) {
 	c.state.containers.Add(ctr.ID(), ctr)
 }
 
+// AddInfraContainer adds a container to the container state store
+func (c *ContainerServer) AddInfraContainer(ctr *oci.Container) {
+	c.stateLock.Lock()
+	defer c.stateLock.Unlock()
+	c.state.infraContainers.Add(ctr.ID(), ctr)
+}
+
 // GetContainer returns a container by its ID
 func (c *ContainerServer) GetContainer(id string) *oci.Container {
 	c.stateLock.Lock()
@@ -599,6 +608,13 @@ func (c *ContainerServer) GetContainer(id string) *oci.Container {
 	return c.state.containers.Get(id)
 }
 
+// GetInfraContainer returns a container by its ID
+func (c *ContainerServer) GetInfraContainer(id string) *oci.Container {
+	c.stateLock.Lock()
+	defer c.stateLock.Unlock()
+	return c.state.infraContainers.Get(id)
+}
+
 // HasContainer checks if a container exists in the state
 func (c *ContainerServer) HasContainer(id string) bool {
 	c.stateLock.Lock()
@@ -617,6 +633,13 @@ func (c *ContainerServer) RemoveContainer(ctr *oci.Container) {
 	c.state.containers.Delete(ctr.ID())
 }
 
+// RemoveInfraContainer removes a container from the container state store
+func (c *ContainerServer) RemoveInfraContainer(ctr *oci.Container) {
+	c.stateLock.Lock()
+	defer c.stateLock.Unlock()
+	c.state.infraContainers.Delete(ctr.ID())
+}
+
 // listContainers returns a list of all containers stored by the server state
 func (c *ContainerServer) listContainers() []*oci.Container {
 	c.stateLock.Lock()
diff --git a/server/server.go b/server/server.go
index a5699b2c..11bc7654 100644
--- a/server/server.go
+++ b/server/server.go
@@ -277,10 +277,18 @@ func (s *Server) addContainer(c *oci.Container) {
 	s.ContainerServer.AddContainer(c)
 }
 
+func (s *Server) addInfraContainer(c *oci.Container) {
+	s.ContainerServer.AddInfraContainer(c)
+}
+
 func (s *Server) getContainer(id string) *oci.Container {
 	return s.ContainerServer.GetContainer(id)
 }
 
+func (s *Server) getInfraContainer(id string) *oci.Container {
+	return s.ContainerServer.GetInfraContainer(id)
+}
+
 // GetSandboxContainer returns the infra container for a given sandbox
 func (s *Server) GetSandboxContainer(id string) *oci.Container {
 	return s.ContainerServer.GetSandboxContainer(id)
@@ -295,6 +303,10 @@ func (s *Server) removeContainer(c *oci.Container) {
 	s.ContainerServer.RemoveContainer(c)
 }
 
+func (s *Server) removeInfraContainer(c *oci.Container) {
+	s.ContainerServer.RemoveInfraContainer(c)
+}
+
 func (s *Server) getPodSandboxFromRequest(podSandboxID string) (*sandbox.Sandbox, error) {
 	if podSandboxID == "" {
 		return nil, sandbox.ErrIDEmpty