From 30748d7ba3f1e5594074e6c8f683eb736b0911d4 Mon Sep 17 00:00:00 2001
From: Raito Bezarius <masterancpp@gmail.com>
Date: Sun, 13 Oct 2024 22:23:31 +0200
Subject: [PATCH] feat(ircv3): add client-batches feature

Signed-off-by: Raito Bezarius <masterancpp@gmail.com>
---
 Cargo.toml        |  1 +
 src/client/mod.rs | 21 ++++++++++++++++++++-
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/Cargo.toml b/Cargo.toml
index 7bb7a5a..5cf45a2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -71,6 +71,7 @@ tokio-rustls = { version = "0.24.0", features = ["dangerous_configuration"], opt
 rustls-pemfile = { version = "1.0.2", optional = true }
 tokio-native-tls = { version = "0.3.1", optional = true }
 webpki-roots = { version = "0.23.0", optional = true }
+atomic-counter = "1.0.1"
 
 
 [dev-dependencies]
diff --git a/src/client/mod.rs b/src/client/mod.rs
index 74c06bc..2a8208a 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -47,6 +47,7 @@
 //! # }
 //! ```
 
+use atomic_counter::{AtomicCounter, RelaxedCounter};
 #[cfg(feature = "ctcp")]
 use chrono::prelude::*;
 use futures_util::{
@@ -82,7 +83,7 @@ use crate::{
         Capability, ChannelMode, Command,
         Command::{
             ChannelMODE, AUTHENTICATE, CAP, INVITE, JOIN, KICK, KILL, NICK, NICKSERV, NOTICE, OPER,
-            PART, PASS, PONG, PRIVMSG, QUIT, SAMODE, SANICK, TOPIC, USER,
+            PART, PASS, PONG, PRIVMSG, QUIT, SAMODE, SANICK, TOPIC, USER, BATCH
         },
         Message, Mode, NegotiationVersion, Response,
     },
@@ -497,6 +498,8 @@ struct ClientState {
     sender: Sender,
     /// The configuration used with this connection.
     config: Config,
+    /// A thread-safe counter for batch IDs
+    batch_id: RelaxedCounter,
     /// A thread-safe map of channels to the list of users in them.
     chanlists: RwLock<HashMap<String, Vec<User>>>,
     /// A thread-safe map of in-progress batch
@@ -512,6 +515,7 @@ impl ClientState {
         ClientState {
             sender,
             config,
+            batch_id: RelaxedCounter::new(0),
             inflight_batches: RwLock::new(HashMap::new()),
             chanlists: RwLock::new(HashMap::new()),
             alt_nick_index: RwLock::new(0),
@@ -870,6 +874,21 @@ impl ClientState {
     }
 
     pub_state_base!();
+
+    /// Sends a client batch with an iterator as the body.
+    pub fn send_client_batch<I: Iterator<Item=Message>>(&mut self, msg_iterator: I) -> error::Result<()> {
+        let batch_id = format!("{}", self.batch_id.get());
+        self.send(BATCH(batch_id.clone(), None, None))?;
+        // Attach to all message that batch ID.
+        msg_iterator.for_each(|msg| {
+            self.send(msg);
+        });
+        // Close the batch.
+        self.send(BATCH(batch_id, None, None))?;
+        self.batch_id.inc();
+
+        Ok(())
+    }
 }
 
 /// Thread-safe sender that can be used with the client.