fix

Author: forsaken628

Cargo.lock

@@ -117,6 +117,7 @@ databend-storages-common-stage = { workspace = true }
 databend-storages-common-table-meta = { workspace = true }
 derive-visitor = { workspace = true }
 dyn-clone = { workspace = true }
+either = { workspace = true }
 enum-as-inner = { workspace = true }
 ethnum = { workspace = true }
 fastrace = { workspace = true }

src/query/service/Cargo.toml

@@ -31,6 +31,7 @@ use databend_common_expression::DataBlock;
 use databend_common_io::prelude::FormatSettings;
 use databend_common_pipeline_transforms::traits::DataBlockSpill;
 use databend_common_pipeline_transforms::traits::Location;
+use either::Either;
 use log::debug;
 use log::info;
 
@@ -160,12 +161,12 @@ impl SizedChannelBuffer {
         self.is_recv_stopped = true
     }
 
-    fn take_block_once(&mut self, builder: &mut PageBuilder) -> result::Result<(), Location> {
+    fn take_block_once(&mut self, builder: &mut PageBuilder) -> Either<(), Location> {
         let Some(block) = self.values.front_mut() else {
-            return Ok(());
+            return Either::Left(());
         };
         let Some(data) = &block.data else {
-            return Err(block.location.clone().unwrap());
+            return Either::Right(block.location.clone().unwrap());
         };
 
         let take_rows = builder.calculate_take_rows(data);
@@ -177,14 +178,16 @@ impl SizedChannelBuffer {
             self.values.pop_front();
             builder.append_full_block(data);
         }
-        Ok(())
+        Either::Left(())
     }
 
-    fn take_block(&mut self, builder: &mut PageBuilder) -> result::Result<bool, Location> {
+    fn take_block(&mut self, builder: &mut PageBuilder) -> Either<bool, Location> {
         while builder.has_capacity() && !self.values.is_empty() {
-            self.take_block_once(builder)?;
+            if let Either::Right(location) = self.take_block_once(builder) {
+                return Either::Right(location);
+            }
         }
-        Ok(!self.values.is_empty())
+        Either::Left(!self.values.is_empty())
     }
 
     fn restore_first(&mut self, location: &Location, data: DataBlock) {
@@ -219,8 +222,8 @@ where S: DataBlockSpill
     #[fastrace::trace(name = "SizedChannel::try_take_block")]
     async fn try_take_block(&self, builder: &mut PageBuilder) -> Result<bool> {
         let location = match self.inner.lock().unwrap().take_block(builder) {
-            Err(location) => location.clone(),
-            Ok(has_more) => {
+            Either::Right(location) => location,
+            Either::Left(has_more) => {
                 return Ok(has_more);
             }
         };

src/query/service/src/servers/http/v1/query/sized_spsc.rs

@@ -171,6 +171,11 @@ ADDRESS = 'https://databend.com';"
         self
     }
 
+    pub fn off_log(mut self) -> ConfigBuilder {
+        self.conf.log.file.on = false;
+        self
+    }
+
     pub fn build(self) -> InnerConfig {
         self.conf
     }

src/query/service/src/test_kits/config.rs

@@ -524,9 +524,9 @@ async fn test_client_compatible_query_id() -> Result<()> {
 
 #[tokio::test(flavor = "current_thread")]
 async fn test_active_sessions() -> Result<()> {
-    let max_active_sessions = 2;
     let conf = ConfigBuilder::create()
-        .max_active_sessions(max_active_sessions)
+        .max_active_sessions(2)
+        .off_log()
         .build();
     let _fixture = TestFixture::setup_with_config(&conf).await?;
     let ep = create_endpoint()?;
@@ -542,6 +542,7 @@ async fn test_active_sessions() -> Result<()> {
         .into_iter()
         .map(|(_status, resp)| (resp.error.map(|e| e.message).unwrap_or_default()))
         .collect::<Vec<_>>();
+    _fixture.keep_alive();
     results.sort();
     let msg = "[HTTP-QUERY] Failed to upgrade session: Current active sessions (2) has exceeded the max_active_sessions limit (2)";
     let expect = vec!["", "", msg];
@@ -1764,12 +1765,12 @@ async fn test_max_size_per_page() -> Result<()> {
     let json = serde_json::json!({"sql": sql.to_string(), "pagination": {"wait_time_secs": wait_time_secs}});
     let (_, reply, body) = TestHttpQueryRequest::new(json).fetch_begin().await?;
     assert!(reply.error.is_none(), "{:?}", reply.error);
-    let len = body.len() as i32;
-    let target = 20_080_000;
-    assert!(len > target);
-    assert!(len < target + 2000);
-    assert_eq!(reply.data.len(), 10_000);
-    assert_eq!(reply.data[0].len(), 2);
+    let target = (10_usize * 1024 * 1024) as f64;
+    assert!(
+        (0.9..1.1).contains(&(body.len() as f64 / target)),
+        "body len {}",
+        body.len()
+    );
     Ok(())
 }