From 39f4db451295dbd8b30db4f94f220182c2c65be9 Mon Sep 17 00:00:00 2001
From: Lucas <nkcs_lykx@hotmail.com>
Date: Thu, 30 Mar 2023 19:08:45 +0800
Subject: [PATCH] Expose the setting of the maximum size of
 `prefill-for-recycle` to users (#301)

Signed-off-by: Lucasliang <nkcs_lykx@hotmail.com>
---
 src/config.rs                     | 57 ++++++++++++++++++++++++++++++-
 src/file_pipe_log/pipe_builder.rs | 13 +++----
 2 files changed, 61 insertions(+), 9 deletions(-)

diff --git a/src/config.rs b/src/config.rs
index 156ae96e..184599d3 100644
--- a/src/config.rs
+++ b/src/config.rs
@@ -1,6 +1,6 @@
 // Copyright (c) 2017-present, PingCAP, Inc. Licensed under Apache-2.0.
 
-use log::warn;
+use log::{info, warn};
 use serde::{Deserialize, Serialize};
 
 use crate::pipe_log::Version;
@@ -108,6 +108,13 @@ pub struct Config {
     ///
     /// Default: false
     pub prefill_for_recycle: bool,
+
+    /// Maximum capacity for preparing log files for recycling when start.
+    /// If not `None`, its size is equal to `purge-threshold`.
+    /// Only available for `prefill-for-recycle` is true.
+    ///
+    /// Default: None
+    pub prefill_limit: Option<ReadableSize>,
 }
 
 impl Default for Config {
@@ -129,6 +136,7 @@ impl Default for Config {
             memory_limit: None,
             enable_log_recycle: false,
             prefill_for_recycle: false,
+            prefill_limit: None,
         };
         // Test-specific configurations.
         #[cfg(test)]
@@ -180,6 +188,14 @@ impl Config {
                 "prefill is not allowed when log recycle is disabled"
             ));
         }
+        if !self.prefill_for_recycle && self.prefill_limit.is_some() {
+            warn!("prefill-limit will be ignored when prefill is disabled");
+            self.prefill_limit = None;
+        }
+        if self.prefill_for_recycle && self.prefill_limit.is_none() {
+            info!("prefill-limit will be calibrated to purge-threshold");
+            self.prefill_limit = Some(self.purge_threshold);
+        }
         #[cfg(not(feature = "swap"))]
         if self.memory_limit.is_some() {
             warn!("memory-limit will be ignored because swap feature is disabled");
@@ -207,6 +223,25 @@ impl Config {
             0
         }
     }
+
+    /// Returns the capacity for preparing log files for recycling when start.
+    pub(crate) fn prefill_capacity(&self) -> usize {
+        // Attention please, log files with Version::V1 could not be recycled, so it's
+        // useless for prefill.
+        if !self.enable_log_recycle || !self.format_version.has_log_signing() {
+            return 0;
+        }
+        let prefill_limit = self.prefill_limit.unwrap_or(ReadableSize(0)).0;
+        if self.prefill_for_recycle && prefill_limit >= self.target_file_size.0 {
+            // Keep same with the maximum setting of `recycle_capacity`.
+            std::cmp::min(
+                (prefill_limit / self.target_file_size.0) as usize + 2,
+                u32::MAX as usize,
+            )
+        } else {
+            0
+        }
+    }
 }
 
 #[cfg(test)]
@@ -304,4 +339,24 @@ mod tests {
             .unwrap()
             .contains("tolerate-corrupted-tail-records"));
     }
+
+    #[test]
+    fn test_prefill_for_recycle() {
+        let default_prefill_v1 = r#"
+            enable-log-recycle = true
+            prefill-for-recycle = true
+        "#;
+        let mut cfg_load: Config = toml::from_str(default_prefill_v1).unwrap();
+        assert!(cfg_load.sanitize().is_ok());
+        assert_eq!(cfg_load.prefill_limit.unwrap(), cfg_load.purge_threshold);
+
+        let default_prefill_v2 = r#"
+            enable-log-recycle = true
+            prefill-for-recycle = false
+            prefill-limit = "20GB"
+        "#;
+        let mut cfg_load: Config = toml::from_str(default_prefill_v2).unwrap();
+        assert!(cfg_load.sanitize().is_ok());
+        assert!(cfg_load.prefill_limit.is_none());
+    }
 }
diff --git a/src/file_pipe_log/pipe_builder.rs b/src/file_pipe_log/pipe_builder.rs
index 57ad95c8..1e3c8e10 100644
--- a/src/file_pipe_log/pipe_builder.rs
+++ b/src/file_pipe_log/pipe_builder.rs
@@ -2,7 +2,6 @@
 
 //! Helper types to recover in-memory states from log files.
 
-use std::cmp;
 use std::fs::{self, File as StdFile};
 use std::io::Write;
 use std::marker::PhantomData;
@@ -32,8 +31,8 @@ use super::pipe::{
 };
 use super::reader::LogItemBatchFileReader;
 
+/// Maximum size for the buffer for prefilling.
 const PREFILL_BUFFER_SIZE: usize = ReadableSize::mb(16).0 as usize;
-const MAX_PREFILL_SIZE: usize = ReadableSize::gb(12).0 as usize;
 
 /// `ReplayMachine` is a type of deterministic state machine that obeys
 /// associative law.
@@ -477,14 +476,12 @@ impl<F: FileSystem> DualPipesBuilder<F> {
 
     fn initialize_files(&mut self) -> Result<()> {
         let target_file_size = self.cfg.target_file_size.0 as usize;
-        let mut target = if self.cfg.prefill_for_recycle {
+        let mut target = std::cmp::min(
+            self.cfg.prefill_capacity(),
             self.cfg
                 .recycle_capacity()
-                .saturating_sub(self.append_files.len())
-        } else {
-            0
-        };
-        target = cmp::min(target, MAX_PREFILL_SIZE / target_file_size);
+                .saturating_sub(self.append_files.len()),
+        );
         let to_create = target.saturating_sub(self.recycled_files.len());
         if to_create > 0 {
             let now = Instant::now();