Allow specifying guid to speed up pool lookup (openzfs#323)

cmd/zfs_object_agent/Cargo.toml

@@ -22,6 +22,7 @@ bincode = "1.3.2"
 bytes = "1.0"
 futures = "0.3.13"
 futures-core = "0.3.13"
+http = "0.2.4"
 lazy_static = "1.4.0"
 log = "0.4"
 lru = "0.6.5"

cmd/zfs_object_agent/src/object_access.rs

@@ -3,6 +3,7 @@ use async_stream::stream;
 use bytes::Bytes;
 use core::time::Duration;
 use futures::{Future, StreamExt};
+use http::StatusCode;
 use lazy_static::lazy_static;
 use log::*;
 use lru::LruCache;
@@ -263,6 +264,28 @@ impl ObjectAccess {
         results
     }
 
+    pub async fn object_exists(&self, key: &str) -> bool {
+        let res = retry(&format!("head {}", prefixed(key)), || async {
+            let req = HeadObjectRequest {
+                bucket: self.bucket_str.clone(),
+                key: prefixed(key),
+                ..Default::default()
+            };
+            let res = self.client.head_object(req).await;
+            match res {
+                Err(RusotoError::Service(HeadObjectError::NoSuchKey(_))) => (false, res),
+                Err(RusotoError::Unknown(rusoto_core::request::BufferedHttpResponse {
+                    status: StatusCode::NOT_FOUND,
+                    body: _,
+                    headers: _,
+                })) => (false, res),
+                _ => (true, res),
+            }
+        })
+        .await;
+        res.err() == None
+    }
+
     async fn put_object_impl(&self, key: &str, data: Vec<u8>) {
         let len = data.len();
         let a = Arc::new(Bytes::from(data));
@@ -353,22 +376,4 @@ impl ObjectAccess {
     pub async fn delete_objects(&self, keys: Vec<String>) {
         while self.delete_objects_impl(&keys).await {}
     }
-
-    pub async fn object_exists(&self, key: &str) -> bool {
-        let prefixed_key = prefixed(key);
-        debug!("looking for {}", prefixed_key);
-        let results = self.list_objects(key, None, None).await;
-
-        assert_eq!(results.len(), 1);
-        let list = &results[0];
-        // Note need to check if this exact name is in the results. If we are looking
-        // for "x/y" and there is "x/y" and "x/yz", both will be returned.
-        list.contents.as_ref().unwrap_or(&vec![]).iter().any(|o| {
-            if let Some(k) = &o.key {
-                k == key
-            } else {
-                false
-            }
-        })
-    }
 }

cmd/zfs_object_agent/src/pool.rs

@@ -142,6 +142,10 @@ impl PoolPhys {
         format!("zfs/{}/super", guid)
     }
 
+    async fn exists(object_access: &ObjectAccess, guid: PoolGUID) -> bool {
+        object_access.object_exists(&Self::key(guid)).await
+    }
+
     async fn get(object_access: &ObjectAccess, guid: PoolGUID) -> Result<Self> {
         let key = Self::key(guid);
         let buf = object_access.get_object(&key).await?;
@@ -390,6 +394,10 @@ impl PoolSyncingState {
 }
 
 impl Pool {
+    pub async fn exists(object_access: &ObjectAccess, guid: PoolGUID) -> bool {
+        PoolPhys::exists(object_access, guid).await
+    }
+
     pub async fn get_config(object_access: &ObjectAccess, guid: PoolGUID) -> Result<NvList> {
         let pool_phys = PoolPhys::get(object_access, guid).await?;
         let ubphys = UberblockPhys::get(object_access, pool_phys.guid, pool_phys.last_txg).await?;

cmd/zfs_object_agent/src/server.rs

@@ -261,6 +261,22 @@ impl Server {
 
         for buck in buckets {
             let object_access = ObjectAccess::from_client(client, buck.as_str());
+            if let Ok(guid) = nvl.lookup_uint64("guid") {
+                if !Pool::exists(&object_access, PoolGUID(guid)).await {
+                    client = object_access.release_client();
+                    continue;
+                }
+                let pool_config = Pool::get_config(&object_access, PoolGUID(guid)).await;
+                if let Err(_) = pool_config {
+                    client = object_access.release_client();
+                    continue;
+                }
+                resp.insert(format!("{}", guid), pool_config.unwrap().as_ref())
+                    .unwrap();
+                debug!("sending response: {:?}", resp);
+                Self::send_response(&self.output, resp).await;
+                return;
+            }
             let objs = object_access
                 .list_objects("zfs/", Some("/".to_string()), None)
                 .await;

lib/libzutil/os/linux/zutil_import_os.c

@@ -437,6 +437,8 @@ zpool_find_import_agent(libpc_handle_t *hdl, importargs_t *iarg,
 	fnvlist_add_string(msg, "region", region);
 	fnvlist_add_string(msg, "endpoint", endpoint);
 	fnvlist_add_string(msg, "credentials", credentials);
+	if (iarg->guid != 0)
+		fnvlist_add_uint64(msg, "guid", iarg->guid);
 	int sock = socket(AF_UNIX, SOCK_STREAM, 0);
 	int err = connect(sock, (struct sockaddr *)&zfs_user_socket,
 	    sizeof (zfs_user_socket));