From 6e73839082ed125640cf5c419377d24b874ef7fd Mon Sep 17 00:00:00 2001
From: Katia <katia.fredes.a@gmail.com>
Date: Sat, 21 Sep 2024 22:24:02 +0900
Subject: [PATCH] add: work counter for server receives maximum limit

---
 src/async_resolver/server_entry.rs | 21 ++++++++-------------
 src/async_resolver/state_block.rs  |  4 ++--
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/src/async_resolver/server_entry.rs b/src/async_resolver/server_entry.rs
index 448884e2..57d7c8f5 100644
--- a/src/async_resolver/server_entry.rs
+++ b/src/async_resolver/server_entry.rs
@@ -4,31 +4,26 @@ use crate::async_resolver::server_info::ServerInfo;
 #[derive(Clone)]
 pub struct ServerEntry {
     info: ServerInfo,
-    retransmissions: u32,
+    work_counter: u16,
 }
 
 impl ServerEntry {
-    pub fn new(info: ServerInfo) -> ServerEntry {
+    pub fn new(info: ServerInfo, work_counter: u16) -> ServerEntry {
         ServerEntry {
             info,
-            retransmissions: 0,
-            // is_active: true,
+            work_counter: work_counter,
         }
     }
 
     pub fn get_info(&self) -> &ServerInfo {
-        &self.info// TODO: see if this is necessary to use clone or not, in order to reuse TCP connections
+        &self.info // TODO: see if this is necessary to use clone or not, in order to reuse TCP connections
     }
 
-    pub fn get_retransmissions(&self) -> u32 {
-        self.retransmissions
+    pub fn get_work_counter(&self) -> u16 {
+        self.work_counter
     }
 
-    pub fn increment_retransmissions(&mut self) {
-        self.retransmissions += 1;
-    }
-
-    pub fn reset_retransmissions(&mut self) {
-        self.retransmissions = 0;
+    pub fn decrement_work_counter(&mut self) {
+        self.work_counter -= 1;
     }
 }
diff --git a/src/async_resolver/state_block.rs b/src/async_resolver/state_block.rs
index 2acc648c..519195d6 100644
--- a/src/async_resolver/state_block.rs
+++ b/src/async_resolver/state_block.rs
@@ -47,11 +47,11 @@ impl StateBlock {
     /// ```
     /// let state_block = StateBlock::new(Instant::now());
     /// ```
-    pub fn new(request_global_limit: u16, servers: Vec<ServerInfo>) -> StateBlock {
+    pub fn new(request_global_limit: u16, server_transmission_limit: u16, servers: Vec<ServerInfo>) -> StateBlock {
         StateBlock {
             timestamp: Instant::now(),
             work_counter: request_global_limit,
-            servers: servers.into_iter().map(ServerEntry::new).collect(),
+            servers: servers.into_iter().map(|server| ServerEntry::new(server, server_transmission_limit)).collect(),
             current_server_index: 0,
         }
     }