risingwave_batch/spill/
spill_op.rs

1// Copyright 2024 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::hash::BuildHasher;
16use std::ops::{Deref, DerefMut};
17use std::path::{Path, PathBuf};
18use std::sync::{Arc, LazyLock};
19
20use anyhow::anyhow;
21use futures_async_stream::try_stream;
22use futures_util::AsyncReadExt;
23use opendal::Operator;
24use opendal::layers::RetryLayer;
25use opendal::services::{Fs, Memory};
26use risingwave_common::array::DataChunk;
27use risingwave_common::util::batch_spill_config::batch_spill_base_dir;
28use risingwave_pb::Message;
29use risingwave_pb::data::DataChunk as PbDataChunk;
30use thiserror_ext::AsReport;
31use tokio::sync::Mutex;
32use twox_hash::XxHash64;
33
34use crate::error::{BatchError, Result};
35use crate::monitor::BatchSpillMetrics;
36
37pub const DEFAULT_SPILL_PARTITION_NUM: usize = 20;
38const RW_MANAGED_SPILL_DIR: &str = "rw_batch_spill/";
39const DEFAULT_IO_BUFFER_SIZE: usize = 256 * 1024;
40const DEFAULT_IO_CONCURRENT_TASK: usize = 8;
41
42#[derive(Clone)]
43pub enum SpillBackend {
44    Disk,
45    /// Only for testing purpose
46    Memory,
47}
48
49/// `SpillOp` is used to manage the spill directory of the spilling executor and it will drop the directory with a RAII style.
50pub struct SpillOp {
51    pub op: Operator,
52}
53
54impl SpillOp {
55    fn batch_spill_root() -> PathBuf {
56        batch_spill_base_dir().join(RW_MANAGED_SPILL_DIR)
57    }
58
59    pub fn create(path: impl AsRef<Path>, spill_backend: SpillBackend) -> Result<SpillOp> {
60        let path = path.as_ref();
61        if !path.is_relative() {
62            bail!("Spill path must be relative, but got {:?}", path);
63        }
64
65        let root = Self::batch_spill_root().join(path);
66
67        let op = match spill_backend {
68            SpillBackend::Disk => {
69                let builder = Fs::default().root(&root.to_string_lossy());
70                Operator::new(builder)?.layer(RetryLayer::default())
71            }
72            SpillBackend::Memory => {
73                let builder = Memory::default().root(&root.to_string_lossy());
74                Operator::new(builder)?.layer(RetryLayer::default())
75            }
76        };
77        Ok(SpillOp { op })
78    }
79
80    pub async fn clean_spill_directory() -> opendal::Result<()> {
81        static LOCK: LazyLock<Mutex<usize>> = LazyLock::new(|| Mutex::new(0));
82        let _guard = LOCK.lock().await;
83
84        let root = Self::batch_spill_root();
85
86        let builder = Fs::default().root(&root.to_string_lossy());
87
88        let op: Operator = Operator::new(builder)?.layer(RetryLayer::default());
89
90        op.delete_with("/").recursive(true).await
91    }
92
93    pub async fn writer_with(&self, name: &str) -> Result<opendal::Writer> {
94        Ok(self
95            .op
96            .writer_with(name)
97            .concurrent(DEFAULT_IO_CONCURRENT_TASK)
98            .chunk(DEFAULT_IO_BUFFER_SIZE)
99            .await?)
100    }
101
102    pub async fn reader_with(&self, name: &str) -> Result<opendal::Reader> {
103        Ok(self
104            .op
105            .reader_with(name)
106            .chunk(DEFAULT_IO_BUFFER_SIZE)
107            .await?)
108    }
109
110    /// spill file content will look like the below.
111    ///
112    /// ```text
113    /// [proto_len]
114    /// [proto_bytes]
115    /// ...
116    /// [proto_len]
117    /// [proto_bytes]
118    /// ```
119    #[try_stream(boxed, ok = DataChunk, error = BatchError)]
120    pub async fn read_stream(reader: opendal::Reader, spill_metrics: Arc<BatchSpillMetrics>) {
121        let mut reader = reader.into_futures_async_read(..).await?;
122        let mut buf = [0u8; 4];
123        loop {
124            if let Err(err) = reader.read_exact(&mut buf).await {
125                if err.kind() == std::io::ErrorKind::UnexpectedEof {
126                    break;
127                } else {
128                    return Err(anyhow!(err).into());
129                }
130            }
131            let len = u32::from_le_bytes(buf) as usize;
132            spill_metrics.batch_spill_read_bytes.inc_by(len as u64 + 4);
133            let mut buf = vec![0u8; len];
134            reader.read_exact(&mut buf).await.map_err(|e| anyhow!(e))?;
135            let chunk_pb: PbDataChunk = Message::decode(buf.as_slice()).map_err(|e| anyhow!(e))?;
136            let chunk = DataChunk::from_protobuf(&chunk_pb)?;
137            yield chunk;
138        }
139    }
140}
141
142impl Drop for SpillOp {
143    fn drop(&mut self) {
144        let op = self.op.clone();
145        tokio::task::spawn(async move {
146            let result = op.delete_with("/").recursive(true).await;
147            if let Err(error) = result {
148                error!(
149                    error = %error.as_report(),
150                    "Failed to remove spill directory"
151                );
152            }
153        });
154    }
155}
156
157impl DerefMut for SpillOp {
158    fn deref_mut(&mut self) -> &mut Self::Target {
159        &mut self.op
160    }
161}
162
163impl Deref for SpillOp {
164    type Target = Operator;
165
166    fn deref(&self) -> &Self::Target {
167        &self.op
168    }
169}
170
171#[derive(Default, Clone, Copy)]
172pub struct SpillBuildHasher(pub u64);
173
174impl BuildHasher for SpillBuildHasher {
175    type Hasher = XxHash64;
176
177    fn build_hasher(&self) -> Self::Hasher {
178        XxHash64::with_seed(self.0)
179    }
180}
181
182pub const SPILL_AT_LEAST_MEMORY: u64 = 1024 * 1024;
risingwave_batch/spill/spill_op.rs

risingwave_batch/spill/
spill_op.rs
