Blob store separation (#1644)

* Add support for JSON and SSE responses to invoke endpoint

* blob store separation

* fix gc url

* fix gc bug

* fix gc test

* initialize blob storage for each namespace

Author: diptanu

server/src/blob_store/mod.rs

@@ -23,6 +23,8 @@ use url::Url;
 
 use crate::metrics::{blob_storage, Timer};
 
+pub mod registry;
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct BlobStorageConfig {
     pub path: String,

server/src/blob_store/registry.rs

@@ -0,0 +1,45 @@
+use std::{
+    collections::HashMap,
+    sync::{Arc, Mutex},
+};
+
+use anyhow::Result;
+
+use super::{BlobStorage, BlobStorageConfig};
+
+pub struct BlobStorageRegistry {
+    pub blob_storage_buckets: Mutex<HashMap<String, Arc<BlobStorage>>>,
+    pub default_blob_storage: Arc<BlobStorage>,
+}
+
+impl BlobStorageRegistry {
+    pub fn new(default_blob_storage_path: &str) -> Result<Self> {
+        let default_blob_storage = BlobStorage::new(BlobStorageConfig {
+            path: default_blob_storage_path.to_string(),
+        })?;
+        Ok(Self {
+            blob_storage_buckets: Mutex::new(HashMap::new()),
+            default_blob_storage: Arc::new(default_blob_storage),
+        })
+    }
+
+    pub fn create_new_blob_store(&self, namespace: &str, path: &str) -> Result<()> {
+        let blob_storage = BlobStorage::new(BlobStorageConfig {
+            path: path.to_string(),
+        })?;
+        self.blob_storage_buckets
+            .lock()
+            .unwrap()
+            .insert(namespace.to_string(), Arc::new(blob_storage));
+        Ok(())
+    }
+
+    pub fn get_blob_store(&self, namespace: &str) -> Arc<BlobStorage> {
+        self.blob_storage_buckets
+            .lock()
+            .unwrap()
+            .get(namespace)
+            .cloned()
+            .unwrap_or_else(|| self.default_blob_storage.clone())
+    }
+}

server/src/data_model/mod.rs

@@ -1772,6 +1772,19 @@ pub struct FunctionExecutorResources {
     pub gpu: Option<GPUResources>, // None if no GPU.
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize, Builder)]
+#[builder(build_fn(skip))]
+pub struct GcUrl {
+    pub url: String,
+    pub namespace: String,
+}
+
+impl GcUrl {
+    pub fn key(&self) -> String {
+        format!("{}|{}", self.namespace, self.url)
+    }
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize, Builder)]
 #[builder(build_fn(skip))]
 pub struct FunctionExecutor {
@@ -2176,6 +2189,7 @@ impl Display for StateChange {
 pub struct Namespace {
     pub name: String,
     pub created_at: u64,
+    pub blob_storage_bucket: Option<String>,
 }
 
 #[cfg(test)]

server/src/executor_api.rs

@@ -28,7 +28,7 @@ use tonic::{Request, Response, Status};
 use tracing::{debug, error, info, instrument, trace, warn};
 
 use crate::{
-    blob_store::{self, BlobStorage},
+    blob_store::{self, registry::BlobStorageRegistry},
     data_model::{
         self,
         Allocation,
@@ -415,7 +415,7 @@ impl TryFrom<FunctionExecutorState> for data_model::FunctionExecutor {
 
 fn to_function_executor_diagnostics(
     function_executor_update: &executor_api_pb::FunctionExecutorUpdate,
-    blob_storage: &BlobStorage,
+    blob_storage_registry: &BlobStorageRegistry,
 ) -> Result<FunctionExecutorDiagnostics> {
     let description = function_executor_update
         .description
@@ -441,15 +441,19 @@ fn to_function_executor_diagnostics(
         .graph_version
         .clone()
         .ok_or(anyhow::anyhow!("graph_version is required"))?;
+    let blob_storage_url_schema = blob_storage_registry
+        .get_blob_store(&namespace)
+        .get_url_scheme();
+    let blob_storage_url = blob_storage_registry.get_blob_store(&namespace).get_url();
     let startup_stdout = prepare_data_payload(
         function_executor_update.startup_stdout.clone(),
-        &blob_storage.get_url_scheme(),
-        &blob_storage.get_url(),
+        &blob_storage_url_schema,
+        &blob_storage_url,
     );
     let startup_stderr = prepare_data_payload(
         function_executor_update.startup_stderr.clone(),
-        &blob_storage.get_url_scheme(),
-        &blob_storage.get_url(),
+        &blob_storage_url_schema,
+        &blob_storage_url,
     );
 
     Ok(data_model::FunctionExecutorDiagnostics {
@@ -465,12 +469,12 @@ fn to_function_executor_diagnostics(
 
 fn to_function_executor_diagnostics_vector(
     function_executor_updates: &Vec<executor_api_pb::FunctionExecutorUpdate>,
-    blob_storage: &BlobStorage,
+    blob_storage_registry: &BlobStorageRegistry,
 ) -> Result<Vec<FunctionExecutorDiagnostics>> {
     function_executor_updates
         .iter()
         .map(|function_executor_update| {
-            to_function_executor_diagnostics(function_executor_update, blob_storage)
+            to_function_executor_diagnostics(function_executor_update, blob_storage_registry)
         })
         .collect()
 }
@@ -479,21 +483,21 @@ pub struct ExecutorAPIService {
     indexify_state: Arc<IndexifyState>,
     executor_manager: Arc<ExecutorManager>,
     api_metrics: Arc<api_io_stats::Metrics>,
-    blob_storage: Arc<blob_store::BlobStorage>,
+    blob_storage_registry: Arc<blob_store::registry::BlobStorageRegistry>,
 }
 
 impl ExecutorAPIService {
     pub fn new(
         indexify_state: Arc<IndexifyState>,
         executor_manager: Arc<ExecutorManager>,
         api_metrics: Arc<api_io_stats::Metrics>,
-        blob_storage: Arc<blob_store::BlobStorage>,
+        blob_storage_registry: Arc<blob_store::registry::BlobStorageRegistry>,
     ) -> Self {
         Self {
             indexify_state,
             executor_manager,
             api_metrics,
-            blob_storage,
+            blob_storage_registry,
         }
     }
 
@@ -577,16 +581,21 @@ impl ExecutorAPIService {
 
             let mut payloads = Vec::new();
             let mut encoding_str = String::new();
+            let blob_storage_url_schema = self
+                .blob_storage_registry
+                .get_blob_store(&namespace)
+                .get_url_scheme();
+            let blob_storage_url = self
+                .blob_storage_registry
+                .get_blob_store(&namespace)
+                .get_url();
             for output in task_result.function_outputs.clone() {
                 let url = output
                     .uri
                     .ok_or(Status::invalid_argument("uri is required"))?;
 
-                let path = blob_store_url_to_path(
-                    &url,
-                    &self.blob_storage.get_url_scheme(),
-                    &self.blob_storage.get_url(),
-                );
+                let path =
+                    blob_store_url_to_path(&url, &blob_storage_url_schema, &blob_storage_url);
                 let size = output
                     .size
                     .ok_or(Status::invalid_argument("size is required"))?;
@@ -621,8 +630,8 @@ impl ExecutorAPIService {
             }
             let invocation_error_payload = prepare_data_payload(
                 task_result.invocation_error_output.clone(),
-                &self.blob_storage.get_url_scheme(),
-                &self.blob_storage.get_url(),
+                &blob_storage_url_schema,
+                &blob_storage_url,
             );
 
             let node_output = NodeOutputBuilder::default()
@@ -642,13 +651,13 @@ impl ExecutorAPIService {
             let task_diagnostic = TaskDiagnostics {
                 stdout: prepare_data_payload(
                     task_result.stdout.clone(),
-                    &self.blob_storage.get_url_scheme(),
-                    &self.blob_storage.get_url(),
+                    &blob_storage_url_schema,
+                    &blob_storage_url,
                 ),
                 stderr: prepare_data_payload(
                     task_result.stderr.clone(),
-                    &self.blob_storage.get_url_scheme(),
-                    &self.blob_storage.get_url(),
+                    &blob_storage_url_schema,
+                    &blob_storage_url,
                 ),
             };
             let allocation_key = Allocation::key_from(
@@ -849,7 +858,7 @@ impl ExecutorApi for ExecutorAPIService {
 
         let function_executor_diagnostics = to_function_executor_diagnostics_vector(
             &executor_update.function_executor_updates,
-            &self.blob_storage,
+            &self.blob_storage_registry,
         )
         .map_err(|e| Status::invalid_argument(e.to_string()))?;
 

server/src/executors.rs

@@ -16,6 +16,7 @@ use tokio::{
 use tracing::{debug, error, trace};
 
 use crate::{
+    blob_store::registry::BlobStorageRegistry,
     data_model::{self, ExecutorId, ExecutorMetadata},
     executor_api::{
         blob_store_path_to_url,
@@ -118,20 +119,18 @@ impl ExecutorRuntimeData {
 }
 
 pub struct ExecutorManager {
-    blob_store_url_scheme: String,
-    blob_store_url: String,
     heartbeat_deadline_queue: Mutex<PriorityQueue<ExecutorId, ReverseInstant>>,
     heartbeat_future: Arc<Mutex<DynamicSleepFuture>>,
     heartbeat_deadline_updater: watch::Sender<Instant>,
     indexify_state: Arc<IndexifyState>,
     runtime_data: RwLock<HashMap<ExecutorId, ExecutorRuntimeData>>,
+    blob_store_registry: Arc<BlobStorageRegistry>,
 }
 
 impl ExecutorManager {
     pub async fn new(
         indexify_state: Arc<IndexifyState>,
-        blob_store_url_scheme: String,
-        blob_store_url: String,
+        blob_store_registry: Arc<BlobStorageRegistry>,
     ) -> Arc<Self> {
         let (heartbeat_future, heartbeat_sender) = DynamicSleepFuture::new(
             far_future(),
@@ -149,8 +148,7 @@ impl ExecutorManager {
             heartbeat_deadline_queue: Mutex::new(PriorityQueue::new()),
             heartbeat_deadline_updater: heartbeat_sender,
             heartbeat_future,
-            blob_store_url_scheme,
-            blob_store_url,
+            blob_store_registry,
         };
 
         let em = Arc::new(em);
@@ -387,20 +385,37 @@ impl ExecutorManager {
         let mut function_executors_pb = vec![];
         let mut task_allocations = vec![];
         for desired_state_fe in desired_executor_state.function_executors.iter() {
+            let blob_store_url_schema = self
+                .blob_store_registry
+                .get_blob_store(
+                    &desired_state_fe
+                        .function_executor
+                        .function_executor
+                        .namespace,
+                )
+                .get_url_scheme();
+            let blob_store_url = self
+                .blob_store_registry
+                .get_blob_store(
+                    &desired_state_fe
+                        .function_executor
+                        .function_executor
+                        .namespace,
+                )
+                .get_url();
             let code_payload_pb = DataPayload {
                 uri: Some(blob_store_path_to_url(
                     &desired_state_fe.code_payload.path,
-                    &self.blob_store_url_scheme,
-                    &self.blob_store_url,
+                    &blob_store_url_schema,
+                    &blob_store_url,
                 )),
                 size: Some(desired_state_fe.code_payload.size),
                 sha256_hash: Some(desired_state_fe.code_payload.sha256_hash.clone()),
                 encoding: Some(DataPayloadEncoding::BinaryZip.into()),
                 encoding_version: None,
             };
             let fe = &desired_state_fe.function_executor.function_executor;
-            let fe_output_payload_uri_prefix =
-                format!("{}/function_executors", self.blob_store_url,);
+            let fe_output_payload_uri_prefix = format!("{}/function_executors", blob_store_url,);
             let fe_description_pb = FunctionExecutorDescription {
                 id: Some(fe.id.get().to_string()),
                 namespace: Some(fe.namespace.clone()),
@@ -466,11 +481,19 @@ impl ExecutorManager {
 
     /// Extracts only the computed fields from a data_model::Task
     pub fn extract_computed_fields(&self, task: &data_model::Task) -> anyhow::Result<ComputedTask> {
+        let blob_store_url_schema = self
+            .blob_store_registry
+            .get_blob_store(&task.namespace)
+            .get_url_scheme();
+        let blob_store_url = self
+            .blob_store_registry
+            .get_blob_store(&task.namespace)
+            .get_url();
         let input_payload = DataPayload {
             uri: Some(blob_store_path_to_url(
                 &task.input.path,
-                &self.blob_store_url_scheme,
-                &self.blob_store_url,
+                &blob_store_url_schema,
+                &blob_store_url,
             )),
             size: Some(task.input.size),
             sha256_hash: Some(task.input.sha256_hash.clone()),
@@ -483,8 +506,8 @@ impl ExecutorManager {
         let reducer_input = task.acc_input.clone().map(|input| DataPayload {
             uri: Some(blob_store_path_to_url(
                 &input.path,
-                &self.blob_store_url_scheme,
-                &self.blob_store_url,
+                &blob_store_url_schema,
+                &blob_store_url,
             )),
             size: Some(input.size),
             sha256_hash: Some(input.sha256_hash.clone()),
@@ -499,7 +522,7 @@ impl ExecutorManager {
         // Create output payload URI prefix
         let output_payload_uri_prefix = format!(
             "{}/{}.{}.{}.{}",
-            self.blob_store_url,
+            blob_store_url,
             task.namespace,
             task.compute_graph_name,
             task.compute_fn_name,