add ut

Author: dqhl76

src/query/service/src/history_tables/meta.rs

@@ -87,16 +87,16 @@ impl HistoryMetaHandle {
             self.node_id.clone(),
             Duration::from_secs(3),
         ))
-        .await
-        .map_err(|_e| "acquire semaphore failed from GlobalHistoryLog")?;
+            .await
+            .map_err(|_e| _e.to_string())?;
         if interval == 0 {
             return Ok(Some(acquired_guard));
         }
         if match ThreadTracker::tracking_future(
             self.meta_client
                 .get_kv(&format!("{}/last_timestamp", meta_key)),
         )
-        .await?
+            .await?
         {
             Some(v) => {
                 let last: u64 = serde_json::from_slice(&v.data)?;
@@ -169,3 +169,78 @@ impl HistoryMetaHandle {
         Ok(())
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::history_tables::meta::HistoryMetaHandle;
+    use databend_common_meta_store::MetaStoreProvider;
+    use std::ops::Deref;
+    use std::time::Duration;
+    use tokio::time::timeout;
+
+    #[tokio::test(flavor = "multi_thread")]
+    pub async fn test_history_table_permit_guard() -> databend_common_exception::Result<()> {
+        let meta_store = MetaStoreProvider::new(Default::default())
+            .create_meta_store()
+            .await
+            .unwrap();
+        let meta_client = meta_store.deref().clone();
+
+
+        let node_id = "test_node_123".to_string();
+        let meta_handle = HistoryMetaHandle::new(meta_client, node_id);
+
+        // Test 1: Basic permit acquisition with interval 0 (no rate limiting)
+        let meta_key = "test/history_table/permit_guard";
+        let guard_result = meta_handle.acquire_with_guard(meta_key, 0).await?;
+        assert!(guard_result.is_some(), "Should acquire permit when interval is 0");
+
+        if let Some(guard) = guard_result {
+            // Verify that the guard contains the correct meta key
+            assert_eq!(guard.meta_key, meta_key);
+        }
+
+        // Same meta key, because we set the interval to 0, it should not block
+        let guard_result2 = timeout(Duration::from_secs(1), meta_handle.acquire_with_guard(meta_key, 0)).await;
+        assert!(guard_result2.is_ok(), "Should acquire permit again when interval is 0");
+        if let Some(guard) = guard_result2.unwrap()? {
+            // Verify that the guard contains the correct meta key
+            assert_eq!(guard.meta_key, meta_key);
+        }
+
+        // Test 2: Permit acquisition with interval > 0 (rate limiting)
+        let meta_key_rate_limited = "test/history_table/permit_guard_rate_limited";
+        let interval_seconds = 2;
+
+        // First acquisition should succeed
+        let first_guard_result = meta_handle.acquire_with_guard(meta_key_rate_limited, interval_seconds).await?;
+        assert!(first_guard_result.is_some(), "First permit acquisition should succeed");
+
+        // Drop the first guard to trigger timestamp update
+        drop(first_guard_result);
+
+        // Immediate second acquisition should fail due to rate limiting
+        let second_guard_result = meta_handle.acquire_with_guard(meta_key_rate_limited, interval_seconds).await?;
+        assert!(second_guard_result.is_none(), "Second permit acquisition should fail due to rate limiting");
+
+        // Test 3: Verify permit guard automatically updates timestamp on drop
+        let meta_key_timestamp = "test/history_table/permit_guard_timestamp";
+
+        // Get initial timestamp (should be None)
+        let initial_timestamp = meta_handle.get_u64_from_meta(&format!("{}/last_timestamp", meta_key_timestamp)).await?;
+        assert!(initial_timestamp.is_none(), "Initial timestamp should be None");
+
+        // Acquire permit with guard
+        let guard = meta_handle.acquire_with_guard(meta_key_timestamp, 0).await?;
+        assert!(guard.is_some(), "Should acquire permit");
+
+        // Drop guard to trigger timestamp update
+        drop(guard);
+
+        // Verify timestamp was updated
+        let updated_timestamp = meta_handle.get_u64_from_meta(&format!("{}/last_timestamp", meta_key_timestamp)).await?;
+        assert!(updated_timestamp.is_some(), "Timestamp should be updated after guard drop");
+
+        Ok(())
+    }
+}