Skip to content

Commit 13f0210

Browse files
TaylorMutchTaylorMutch
committed
refactor(kubernetes): collapse agent-sandbox api selection
Signed-off-by: Taylor Mutch <taylormutch@gmail.com>
1 parent 7f9c1ca commit 13f0210

1 file changed

Lines changed: 52 additions & 34 deletions

File tree

  • crates/openshell-driver-kubernetes/src

‎crates/openshell-driver-kubernetes/src/driver.rs‎

Lines changed: 52 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -89,6 +89,11 @@ pub const SANDBOX_KIND: &str = "Sandbox";
8989
const GPU_RESOURCE_NAME: &str = "nvidia.com/gpu";
9090
const SPIFFE_WORKLOAD_API_VOLUME_NAME: &str = "spiffe-workload-api";
9191

92+
struct AgentSandboxApi {
93+
api: Api<DynamicObject>,
94+
resource: ApiResource,
95+
}
96+
9297
// This POC treats the selected Struct as a driver-local typed schema. Once the
9398
// Kubernetes shape stabilizes, these serde structs may move to driver-local
9499
// protobuf definitions, but the typed decode should stay inside this driver.
@@ -261,10 +266,16 @@ impl KubernetesComputeDriver {
261266
&self.config.ssh_socket_path
262267
}
263268

264-
fn sandbox_api(&self, client: Client, sandbox_api_version: &str) -> Api<DynamicObject> {
269+
fn agent_sandbox_api(&self, client: Client, sandbox_api_version: &str) -> AgentSandboxApi {
265270
let gvk = GroupVersionKind::gvk(SANDBOX_GROUP, sandbox_api_version, SANDBOX_KIND);
266271
let resource = ApiResource::from_gvk(&gvk);
267-
Api::namespaced_with(client, &self.config.namespace, &resource)
272+
let api = Api::namespaced_with(client, &self.config.namespace, &resource);
273+
AgentSandboxApi { api, resource }
274+
}
275+
276+
async fn supported_agent_sandbox_api(&self, client: Client) -> Result<AgentSandboxApi, String> {
277+
let sandbox_api_version = self.supported_sandbox_api_version(client.clone()).await?;
278+
Ok(self.agent_sandbox_api(client, sandbox_api_version))
268279
}
269280

270281
async fn supported_sandbox_api_version(&self, client: Client) -> Result<&'static str, String> {
@@ -281,9 +292,12 @@ impl KubernetesComputeDriver {
281292
client: Client,
282293
) -> Result<&'static str, String> {
283294
for sandbox_api_version in SANDBOX_VERSIONS {
284-
let api = self.sandbox_api(client.clone(), sandbox_api_version);
285-
match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(&ListParams::default().limit(1)))
286-
.await
295+
let agent_sandbox_api = self.agent_sandbox_api(client.clone(), sandbox_api_version);
296+
match tokio::time::timeout(
297+
KUBE_API_TIMEOUT,
298+
agent_sandbox_api.api.list(&ListParams::default().limit(1)),
299+
)
300+
.await
287301
{
288302
Ok(Ok(_)) => {
289303
debug!(
@@ -354,11 +368,10 @@ impl KubernetesComputeDriver {
354368
"Fetching sandbox from Kubernetes"
355369
);
356370

357-
let sandbox_api_version = self
358-
.supported_sandbox_api_version(self.client.clone())
371+
let agent_sandbox_api = self
372+
.supported_agent_sandbox_api(self.client.clone())
359373
.await?;
360-
let api = self.sandbox_api(self.client.clone(), sandbox_api_version);
361-
match tokio::time::timeout(KUBE_API_TIMEOUT, api.get(name)).await {
374+
match tokio::time::timeout(KUBE_API_TIMEOUT, agent_sandbox_api.api.get(name)).await {
362375
Ok(Ok(obj)) => sandbox_from_object(&self.config.namespace, obj).map(Some),
363376
Ok(Err(KubeError::Api(err))) if err.code == 404 => {
364377
debug!(sandbox_name = %name, "Sandbox not found in Kubernetes");
@@ -392,11 +405,15 @@ impl KubernetesComputeDriver {
392405
"Listing sandboxes from Kubernetes"
393406
);
394407

395-
let sandbox_api_version = self
396-
.supported_sandbox_api_version(self.client.clone())
408+
let agent_sandbox_api = self
409+
.supported_agent_sandbox_api(self.client.clone())
397410
.await?;
398-
let api = self.sandbox_api(self.client.clone(), sandbox_api_version);
399-
match tokio::time::timeout(KUBE_API_TIMEOUT, api.list(&ListParams::default())).await {
411+
match tokio::time::timeout(
412+
KUBE_API_TIMEOUT,
413+
agent_sandbox_api.api.list(&ListParams::default()),
414+
)
415+
.await
416+
{
400417
Ok(Ok(list)) => {
401418
let mut sandboxes = list
402419
.items
@@ -450,13 +467,11 @@ impl KubernetesComputeDriver {
450467
"Creating sandbox in Kubernetes"
451468
);
452469

453-
let sandbox_api_version = self
454-
.supported_sandbox_api_version(self.client.clone())
470+
let agent_sandbox_api = self
471+
.supported_agent_sandbox_api(self.client.clone())
455472
.await
456473
.map_err(KubernetesDriverError::Message)?;
457-
let gvk = GroupVersionKind::gvk(SANDBOX_GROUP, sandbox_api_version, SANDBOX_KIND);
458-
let resource = ApiResource::from_gvk(&gvk);
459-
let mut obj = DynamicObject::new(name, &resource);
474+
let mut obj = DynamicObject::new(name, &agent_sandbox_api.resource);
460475
obj.metadata = ObjectMeta {
461476
name: Some(name.to_string()),
462477
namespace: Some(self.config.namespace.clone()),
@@ -488,9 +503,11 @@ impl KubernetesComputeDriver {
488503
.provider_spiffe_workload_api_socket_path,
489504
};
490505
obj.data = sandbox_to_k8s_spec(sandbox.spec.as_ref(), &params);
491-
let api = self.sandbox_api(self.client.clone(), sandbox_api_version);
492-
493-
match tokio::time::timeout(KUBE_API_TIMEOUT, api.create(&PostParams::default(), &obj)).await
506+
match tokio::time::timeout(
507+
KUBE_API_TIMEOUT,
508+
agent_sandbox_api.api.create(&PostParams::default(), &obj),
509+
)
510+
.await
494511
{
495512
Ok(Ok(_result)) => {
496513
info!(
@@ -531,12 +548,14 @@ impl KubernetesComputeDriver {
531548
"Deleting sandbox from Kubernetes"
532549
);
533550

534-
let sandbox_api_version = self
535-
.supported_sandbox_api_version(self.client.clone())
551+
let agent_sandbox_api = self
552+
.supported_agent_sandbox_api(self.client.clone())
536553
.await?;
537-
let api = self.sandbox_api(self.client.clone(), sandbox_api_version);
538-
match tokio::time::timeout(KUBE_API_TIMEOUT, api.delete(name, &DeleteParams::default()))
539-
.await
554+
match tokio::time::timeout(
555+
KUBE_API_TIMEOUT,
556+
agent_sandbox_api.api.delete(name, &DeleteParams::default()),
557+
)
558+
.await
540559
{
541560
Ok(Ok(_response)) => {
542561
info!(sandbox_name = %name, "Sandbox deleted from Kubernetes");
@@ -569,11 +588,10 @@ impl KubernetesComputeDriver {
569588
}
570589

571590
pub async fn sandbox_exists(&self, name: &str) -> Result<bool, String> {
572-
let sandbox_api_version = self
573-
.supported_sandbox_api_version(self.client.clone())
591+
let agent_sandbox_api = self
592+
.supported_agent_sandbox_api(self.client.clone())
574593
.await?;
575-
let api = self.sandbox_api(self.client.clone(), sandbox_api_version);
576-
match tokio::time::timeout(KUBE_API_TIMEOUT, api.get(name)).await {
594+
match tokio::time::timeout(KUBE_API_TIMEOUT, agent_sandbox_api.api.get(name)).await {
577595
Ok(Ok(_)) => Ok(true),
578596
Ok(Err(KubeError::Api(err))) if err.code == 404 => Ok(false),
579597
Ok(Err(err)) => Err(err.to_string()),
@@ -588,12 +606,12 @@ impl KubernetesComputeDriver {
588606
#[allow(clippy::unused_async)]
589607
pub async fn watch_sandboxes(&self) -> Result<WatchStream, String> {
590608
let namespace = self.config.namespace.clone();
591-
let sandbox_api_version = self
592-
.supported_sandbox_api_version(self.watch_client.clone())
609+
let agent_sandbox_api = self
610+
.supported_agent_sandbox_api(self.watch_client.clone())
593611
.await?;
594-
let sandbox_api = self.sandbox_api(self.watch_client.clone(), sandbox_api_version);
595612
let event_api: Api<KubeEventObj> = Api::namespaced(self.watch_client.clone(), &namespace);
596-
let mut sandbox_stream = watcher::watcher(sandbox_api, watcher::Config::default()).boxed();
613+
let mut sandbox_stream =
614+
watcher::watcher(agent_sandbox_api.api, watcher::Config::default()).boxed();
597615
let mut event_stream = watcher::watcher(event_api, watcher::Config::default()).boxed();
598616
let (tx, rx) = mpsc::channel(256);
599617

0 commit comments

Comments
 (0)