Merge e199899 into 56b0695

Author: rklaehn

src/rpc.rs

@@ -317,6 +317,16 @@ impl<D: crate::store::Store> Handler<D> {
             #[allow(clippy::manual_flatten)]
             for item in tags {
                 if let Ok((name, HashAndFormat { hash, format })) = item {
+                    if let Some(from) = msg.from.as_ref() {
+                        if &name < from {
+                            continue;
+                        }
+                    }
+                    if let Some(to) = msg.to.as_ref() {
+                        if &name >= to {
+                            break;
+                        }
+                    }
                     if (format.is_raw() && msg.raw) || (format.is_hash_seq() && msg.hash_seq) {
                         co.yield_(TagInfo { name, hash, format }).await;
                     }

src/rpc/client/tags.rs

@@ -42,6 +42,27 @@ where
         Self { rpc }
     }
 
+    /// Get the value of a single tag
+    pub async fn get(&self, name: impl AsRef<[u8]>) -> Result<Option<TagInfo>> {
+        let mut stream = self
+            .rpc
+            .server_streaming(ListRequest::single(name.as_ref()))
+            .await?;
+        Ok(stream.next().await.transpose()?)
+    }
+
+    /// Lists all tags.
+    pub async fn list_prefix(
+        &self,
+        prefix: impl AsRef<[u8]>,
+    ) -> Result<impl Stream<Item = Result<TagInfo>>> {
+        let stream = self
+            .rpc
+            .server_streaming(ListRequest::prefix(prefix.as_ref()))
+            .await?;
+        Ok(stream.map(|res| res.map_err(anyhow::Error::from)))
+    }
+
     /// Lists all tags.
     pub async fn list(&self) -> Result<impl Stream<Item = Result<TagInfo>>> {
         let stream = self.rpc.server_streaming(ListRequest::all()).await?;

src/rpc/proto/tags.rs

@@ -1,10 +1,16 @@
 //! Tags RPC protocol
+use bytes::Bytes;
 use nested_enum_utils::enum_conversions;
 use quic_rpc_derive::rpc_requests;
 use serde::{Deserialize, Serialize};
 
 use super::{RpcResult, RpcService};
-use crate::{net_protocol::BatchId, rpc::client::tags::TagInfo, HashAndFormat, Tag};
+use crate::{
+    net_protocol::BatchId,
+    rpc::client::tags::TagInfo,
+    util::{increment_vec, next_prefix},
+    HashAndFormat, Tag,
+};
 
 #[allow(missing_docs)]
 #[derive(strum::Display, Debug, Serialize, Deserialize)]
@@ -73,14 +79,53 @@ pub struct ListRequest {
     pub raw: bool,
     /// List hash seq tags
     pub hash_seq: bool,
+    /// From tag
+    pub from: Option<Tag>,
+    /// To tag (exclusive)
+    pub to: Option<Tag>,
 }
 
 impl ListRequest {
+    /// List tags with a prefix
+    pub fn prefix(prefix: &[u8]) -> Self {
+        let from = prefix.to_vec();
+        let mut to = from.clone();
+        let from = Bytes::from(from).into();
+        let to = if next_prefix(&mut to) {
+            Some(Bytes::from(to).into())
+        } else {
+            None
+        };
+        Self {
+            raw: true,
+            hash_seq: true,
+            from: Some(from),
+            to,
+        }
+    }
+
+    /// List a single tag
+    pub fn single(name: &[u8]) -> Self {
+        let from = name.to_vec();
+        let mut next = from.clone();
+        increment_vec(&mut next);
+        let from = Bytes::from(from).into();
+        let to = Bytes::from(next).into();
+        Self {
+            raw: true,
+            hash_seq: true,
+            from: Some(from),
+            to: Some(to),
+        }
+    }
+
     /// List all tags
     pub fn all() -> Self {
         Self {
             raw: true,
             hash_seq: true,
+            from: None,
+            to: None,
         }
     }
 
@@ -89,6 +134,8 @@ impl ListRequest {
         Self {
             raw: true,
             hash_seq: false,
+            from: None,
+            to: None,
         }
     }
 
@@ -97,6 +144,8 @@ impl ListRequest {
         Self {
             raw: false,
             hash_seq: true,
+            from: None,
+            to: None,
         }
     }
 }

src/util.rs

@@ -302,6 +302,36 @@ pub(crate) fn raw_outboard_size(size: u64) -> u64 {
     BaoTree::new(size, IROH_BLOCK_SIZE).outboard_size()
 }
 
+/// Given a prefix, increment it lexographically.
+///
+/// If the prefix is all FF, this will return false because there is no
+/// higher prefix than that.
+#[allow(dead_code)]
+pub(crate) fn next_prefix(bytes: &mut [u8]) -> bool {
+    for byte in bytes.iter_mut().rev() {
+        if *byte < 255 {
+            *byte += 1;
+            return true;
+        }
+        *byte = 0;
+    }
+    false
+}
+
+/// Increment a byte vector, lexographically.
+#[allow(dead_code)]
+pub(crate) fn increment_vec(bytes: &mut Vec<u8>) {
+    for byte in bytes.iter_mut().rev() {
+        if *byte < 255 {
+            *byte += 1;
+            return;
+        }
+        *byte = 0;
+    }
+
+    bytes.push(0);
+}
+
 /// Synchronously compute the outboard of a file, and return hash and outboard.
 ///
 /// It is assumed that the file is not modified while this is running.