From 1431593a007208de943ddabcf58b57e6b006dd6d Mon Sep 17 00:00:00 2001
From: Klaus Ma <klausm@nvidia.com>
Date: Fri, 26 Apr 2024 07:45:00 +0000
Subject: [PATCH] Load container/sandbox when start.

Signed-off-by: Klaus Ma <klausm@nvidia.com>
---
 shim/src/cri/runtime/mod.rs | 37 +++++++++++++++++++++++++++++++++++--
 1 file changed, 35 insertions(+), 2 deletions(-)

diff --git a/shim/src/cri/runtime/mod.rs b/shim/src/cri/runtime/mod.rs
index 4361b00..2b18f3c 100644
--- a/shim/src/cri/runtime/mod.rs
+++ b/shim/src/cri/runtime/mod.rs
@@ -43,6 +43,7 @@ use self::crirpc::{
 use self::storage::Storage;
 use crate::cfg::ChariotOptions;
 use crate::common::ChariotError;
+use crate::cri::runtime::storage::{Container, Sandbox};
 use crate::rpc::cri as crirpc;
 
 mod storage;
@@ -56,6 +57,8 @@ pub struct RuntimeShim {
 impl RuntimeShim {
     pub async fn connect(opts: ChariotOptions) -> Result<Self, ChariotError> {
         let mut clients = HashMap::new();
+        let storage = storage::new(&opts.storage)?;
+
         for rt in opts.runtimes {
             let uri = rt.endpoint.find("://").unwrap_or(rt.endpoint.len());
             let (protocol, path) = rt.endpoint.split_at(uri + 3);
@@ -92,11 +95,41 @@ impl RuntimeShim {
                 rt.name, resp.runtime_name, resp.runtime_version
             );
 
+            // Sync up Sandbox and Container of each runtime.
+            let request = crirpc::ListPodSandboxRequest { filter: None };
+            let sandbox = client
+                .list_pod_sandbox(request)
+                .await
+                .map_err(|e| ChariotError::Cri(e.to_string()))?;
+
+            for s in sandbox.into_inner().items {
+                storage
+                    .persist_sandbox(&Sandbox {
+                        id: s.id.clone(),
+                        runtime: rt.name.clone(),
+                    })
+                    .await?;
+            }
+
+            let request = crirpc::ListContainersRequest { filter: None };
+            let container = client
+                .list_containers(request)
+                .await
+                .map_err(|e| ChariotError::Cri(e.to_string()))?;
+
+            for c in container.into_inner().containers {
+                storage
+                    .persist_container(&Container {
+                        id: c.id.clone(),
+                        sandbox: c.pod_sandbox_id.clone(),
+                        runtime: rt.name.clone(),
+                    })
+                    .await?;
+            }
+
             clients.insert(rt.name, client);
         }
 
-        let storage = storage::new(&opts.storage)?;
-
         let c = storage.list_containers().await?;
         let s = storage.list_sandboxes().await?;