feat(tvix/castore): set user-agent for object_store blob/directorysvc r/8968

Change-Id: I9fcebffb19174cba2e9001398419d3041266400c
Reviewed-on: https://cl.tvl.fyi/c/depot/+/12850
Reviewed-by: Vladimir Kryachko <v.kryachko@gmail.com>
Tested-by: BuildkiteCI
Autosubmit: flokli <flokli@flokli.de>
Reviewed-by: flokli <flokli@flokli.de>

3 files changed, 41 insertions, 9 deletions

diff --git a/tvix/castore/src/blobservice/object_store.rs b/tvix/castore/src/blobservice/object_store.rs
index 10874af64011..317e8c18e893 100644
--- a/tvix/castore/src/blobservice/object_store.rs
+++ b/tvix/castore/src/blobservice/object_store.rs
@@ -1,5 +1,5 @@
 use std::{
-    collections::HashMap,
+    collections::{hash_map, HashMap},
     io::{self, Cursor},
     pin::pin,
     sync::Arc,
@@ -298,10 +298,24 @@ impl ServiceBuilder for ObjectStoreBlobServiceConfig {
         instance_name: &str,
         _context: &CompositionContext,
     ) -> Result<Arc<dyn BlobService>, Box<dyn std::error::Error + Send + Sync + 'static>> {
-        let (object_store, path) = object_store::parse_url_opts(
-            &self.object_store_url.parse()?,
-            &self.object_store_options,
-        )?;
+        let opts = {
+            let mut opts: HashMap<&str, _> = self
+                .object_store_options
+                .iter()
+                .map(|(k, v)| (k.as_str(), v.as_str()))
+                .collect();
+
+            if let hash_map::Entry::Vacant(e) =
+                opts.entry(object_store::ClientConfigKey::UserAgent.as_ref())
+            {
+                e.insert(crate::USER_AGENT);
+            }
+
+            opts
+        };
+
+        let (object_store, path) =
+            object_store::parse_url_opts(&self.object_store_url.parse()?, opts)?;
         Ok(Arc::new(ObjectStoreBlobService {
             instance_name: instance_name.to_string(),
             object_store: Arc::new(object_store),
diff --git a/tvix/castore/src/directoryservice/object_store.rs b/tvix/castore/src/directoryservice/object_store.rs
index 1916e59eacaa..147a79cbfb42 100644
--- a/tvix/castore/src/directoryservice/object_store.rs
+++ b/tvix/castore/src/directoryservice/object_store.rs
@@ -1,3 +1,4 @@
+use std::collections::hash_map;
 use std::collections::HashMap;
 use std::sync::Arc;
 
@@ -232,10 +233,24 @@ impl ServiceBuilder for ObjectStoreDirectoryServiceConfig {
         instance_name: &str,
         _context: &CompositionContext,
     ) -> Result<Arc<dyn DirectoryService>, Box<dyn std::error::Error + Send + Sync + 'static>> {
-        let (object_store, path) = object_store::parse_url_opts(
-            &self.object_store_url.parse()?,
-            &self.object_store_options,
-        )?;
+        let opts = {
+            let mut opts: HashMap<&str, _> = self
+                .object_store_options
+                .iter()
+                .map(|(k, v)| (k.as_str(), v.as_str()))
+                .collect();
+
+            if let hash_map::Entry::Vacant(e) =
+                opts.entry(object_store::ClientConfigKey::UserAgent.as_ref())
+            {
+                e.insert(crate::USER_AGENT);
+            }
+
+            opts
+        };
+
+        let (object_store, path) =
+            object_store::parse_url_opts(&self.object_store_url.parse()?, opts)?;
         Ok(Arc::new(ObjectStoreDirectoryService::new(
             instance_name.to_string(),
             Arc::new(object_store),
diff --git a/tvix/castore/src/lib.rs b/tvix/castore/src/lib.rs
index 93d06fd4582d..78018d8dfcbd 100644
--- a/tvix/castore/src/lib.rs
+++ b/tvix/castore/src/lib.rs
@@ -21,6 +21,9 @@ pub mod import;
 pub mod proto;
 pub mod tonic;
 
+// Used as user agent in various HTTP Clients
+const USER_AGENT: &str = concat!(env!("CARGO_PKG_NAME"), "/", env!("CARGO_PKG_VERSION"));
+
 pub use digests::{B3Digest, B3_LEN};
 pub use errors::{DirectoryError, Error, ValidateNodeError};
 pub use hashing_reader::{B3HashingReader, HashingReader};