object_store/layers/
hdfs.rs

1// Copyright 2023 Greptime Team
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::fmt::{self, Debug};
16use std::sync::Arc;
17
18use common_telemetry::warn;
19use hdfs_native::{Client, ClientBuilder};
20use opendal::raw::oio::{Delete as _, Read as _, ReadStream as _, Write as _};
21use opendal::raw::{
22    Layer, OpCompose, OpCopy, OpCreateDir, OpDelete, OpList, OpPresign, OpRead, OpRename, OpStat,
23    OpWrite, RpCreateDir, RpPresign, RpRename, RpStat, Service, ServiceInfo, Servicer, oio,
24};
25use opendal::{Buffer, Capability, ErrorKind, Metadata, OperationContext, Result};
26use uuid::Uuid;
27
28/// Adds atomic writes and streaming copies to the native HDFS backend.
29#[derive(Debug, Clone)]
30pub struct HdfsCompatibilityLayer {
31    renamer: AtomicRenamer,
32}
33
34impl HdfsCompatibilityLayer {
35    /// Creates a compatibility layer for the HDFS connection.
36    pub fn new(
37        name_node: &str,
38        root: &str,
39        options: &std::collections::HashMap<String, String>,
40    ) -> Result<Self> {
41        let mut config = std::collections::HashMap::new();
42        let namenodes = name_node
43            .split(',')
44            .filter_map(|value| {
45                let value = value
46                    .trim()
47                    .trim_start_matches("hdfs://")
48                    .trim_end_matches('/');
49                (!value.is_empty()).then_some(value)
50            })
51            .collect::<Vec<_>>();
52        for (index, namenode) in namenodes.iter().enumerate() {
53            config.insert(
54                format!("dfs.namenode.rpc-address.nameservice.nn{index}"),
55                (*namenode).to_string(),
56            );
57        }
58        config.insert(
59            "dfs.ha.namenodes.nameservice".to_string(),
60            (0..namenodes.len())
61                .map(|index| format!("nn{index}"))
62                .collect::<Vec<_>>()
63                .join(","),
64        );
65        config.extend(options.clone());
66        let client = ClientBuilder::new()
67            .with_url("hdfs://nameservice")
68            .with_config(config)
69            .build()
70            .map_err(hdfs_error)?;
71        Ok(Self {
72            renamer: AtomicRenamer::Native {
73                client,
74                root: opendal::raw::normalize_root(root),
75            },
76        })
77    }
78
79    /// Creates a compatibility layer backed by the inner service's rename.
80    #[cfg(any(test, feature = "testing"))]
81    pub fn new_for_test() -> Self {
82        Self {
83            renamer: AtomicRenamer::Raw,
84        }
85    }
86}
87
88#[derive(Debug, Clone)]
89enum AtomicRenamer {
90    Native {
91        client: Client,
92        root: String,
93    },
94    #[cfg(any(test, feature = "testing"))]
95    Raw,
96}
97
98impl AtomicRenamer {
99    async fn rename(
100        &self,
101        _inner: &Servicer,
102        _ctx: &OperationContext,
103        from: &str,
104        to: &str,
105    ) -> Result<()> {
106        match self {
107            Self::Native { client, root } => {
108                // OpenDAL's HDFS rename removes an existing destination before
109                // renaming. Use HDFS Rename2 with overwrite to keep replacement atomic.
110                client
111                    .rename(
112                        &opendal::raw::build_rooted_abs_path(root, from),
113                        &opendal::raw::build_rooted_abs_path(root, to),
114                        true,
115                    )
116                    .await
117                    .map_err(hdfs_error)
118            }
119            #[cfg(any(test, feature = "testing"))]
120            Self::Raw => _inner
121                .rename(_ctx, from, to, OpRename::new())
122                .await
123                .map(|_| ()),
124        }
125    }
126}
127
128fn hdfs_error(error: hdfs_native::HdfsError) -> opendal::Error {
129    opendal::Error::new(ErrorKind::Unexpected, "native HDFS operation failed").set_source(error)
130}
131
132impl Layer for HdfsCompatibilityLayer {
133    fn apply_service(&self, inner: Servicer) -> Servicer {
134        Arc::new(HdfsCompatibilityService {
135            inner,
136            renamer: self.renamer.clone(),
137        })
138    }
139}
140
141struct HdfsCompatibilityService {
142    inner: Servicer,
143    renamer: AtomicRenamer,
144}
145
146impl Debug for HdfsCompatibilityService {
147    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
148        f.debug_struct("HdfsCompatibilityService")
149            .field("inner", &self.inner)
150            .finish()
151    }
152}
153
154/// A writer that publishes non-append writes with an atomic rename.
155struct HdfsWriter(HdfsWriterInner);
156
157enum HdfsWriterInner {
158    Direct(oio::Writer),
159    Atomic {
160        inner: Servicer,
161        context: OperationContext,
162        renamer: AtomicRenamer,
163        writer: Option<oio::Writer>,
164        temporary_path: String,
165        target_path: String,
166        /// Message of the write that failed and removed the writer, so later
167        /// `write`/`close` calls can explain why the writer is gone.
168        failure: Option<String>,
169    },
170}
171
172impl oio::Write for HdfsWriter {
173    async fn write(&mut self, buffer: Buffer) -> Result<()> {
174        match &mut self.0 {
175            HdfsWriterInner::Direct(writer) => writer.write(buffer).await,
176            HdfsWriterInner::Atomic {
177                inner,
178                context,
179                writer,
180                temporary_path,
181                failure,
182                ..
183            } => {
184                let result = writer
185                    .as_mut()
186                    .ok_or_else(|| writer_unavailable(failure.as_deref()))?
187                    .write(buffer)
188                    .await;
189                if let Err(error) = result {
190                    // A failed write leaves the temporary file behind, so abort
191                    // it before returning the original error to the caller.
192                    if let Some(writer) = writer.take() {
193                        abort_and_delete(inner, context, writer, temporary_path).await;
194                    }
195                    *failure = Some(error.to_string());
196                    return Err(error);
197                }
198                Ok(())
199            }
200        }
201    }
202
203    async fn close(&mut self) -> Result<Metadata> {
204        match &mut self.0 {
205            HdfsWriterInner::Direct(writer) => writer.close().await,
206            HdfsWriterInner::Atomic {
207                inner,
208                context,
209                renamer,
210                writer,
211                temporary_path,
212                target_path,
213                failure,
214            } => {
215                let mut writer = writer
216                    .take()
217                    .ok_or_else(|| writer_unavailable(failure.as_deref()))?;
218                let metadata = match writer.close().await {
219                    Ok(metadata) => metadata,
220                    Err(error) => {
221                        drop(writer);
222                        let _ = delete_path(inner, context, temporary_path).await;
223                        return Err(error);
224                    }
225                };
226                drop(writer);
227
228                if let Err(error) = renamer
229                    .rename(inner, context, temporary_path, target_path)
230                    .await
231                {
232                    let _ = delete_path(inner, context, temporary_path).await;
233                    return Err(error);
234                }
235
236                Ok(metadata)
237            }
238        }
239    }
240
241    async fn abort(&mut self) -> Result<()> {
242        match &mut self.0 {
243            HdfsWriterInner::Direct(writer) => writer.abort().await,
244            HdfsWriterInner::Atomic {
245                inner,
246                context,
247                writer,
248                temporary_path,
249                ..
250            } => {
251                let abort_result = if let Some(mut writer) = writer.take() {
252                    let result = writer.abort().await;
253                    drop(writer);
254                    result
255                } else {
256                    Ok(())
257                };
258                let cleanup_result = delete_path(inner, context, temporary_path).await;
259
260                match abort_result {
261                    Err(error) if error.kind() != ErrorKind::Unsupported => Err(error),
262                    _ => cleanup_result,
263                }
264            }
265        }
266    }
267}
268
269fn writer_unavailable(failure: Option<&str>) -> opendal::Error {
270    match failure {
271        Some(failure) => opendal::Error::new(
272            ErrorKind::Unexpected,
273            format!("HDFS writer is unavailable: a previous write failed: {failure}"),
274        ),
275        None => opendal::Error::new(
276            ErrorKind::Unexpected,
277            "HDFS writer is unavailable after close or abort",
278        ),
279    }
280}
281
282impl Service for HdfsCompatibilityService {
283    type Reader = oio::Reader;
284    type Writer = HdfsWriter;
285    type Lister = oio::Lister;
286    type Deleter = oio::Deleter;
287    type Copier = oio::OneShotCopier;
288    type Composer = oio::Composer;
289
290    fn info(&self) -> ServiceInfo {
291        self.inner.info()
292    }
293
294    fn capability(&self) -> Capability {
295        let mut capability = self.inner.capability();
296        capability.copy = true;
297        capability
298    }
299
300    async fn create_dir(
301        &self,
302        ctx: &OperationContext,
303        path: &str,
304        args: OpCreateDir,
305    ) -> Result<RpCreateDir> {
306        self.inner.create_dir(ctx, path, args).await
307    }
308
309    async fn stat(&self, ctx: &OperationContext, path: &str, args: OpStat) -> Result<RpStat> {
310        self.inner.stat(ctx, path, args).await
311    }
312
313    fn read(&self, ctx: &OperationContext, path: &str, args: OpRead) -> Result<Self::Reader> {
314        self.inner.read(ctx, path, args)
315    }
316
317    fn write(&self, ctx: &OperationContext, path: &str, args: OpWrite) -> Result<Self::Writer> {
318        if args.append() {
319            return self
320                .inner
321                .write(ctx, path, args)
322                .map(|writer| HdfsWriter(HdfsWriterInner::Direct(writer)));
323        }
324
325        let temporary_path = temporary_path(path);
326        let writer = self.inner.write(ctx, &temporary_path, args)?;
327        Ok(HdfsWriter(HdfsWriterInner::Atomic {
328            inner: Arc::clone(&self.inner),
329            context: ctx.clone(),
330            renamer: self.renamer.clone(),
331            writer: Some(writer),
332            temporary_path,
333            target_path: path.to_string(),
334            failure: None,
335        }))
336    }
337
338    fn copy(
339        &self,
340        ctx: &OperationContext,
341        from: &str,
342        to: &str,
343        args: OpCopy,
344    ) -> Result<Self::Copier> {
345        if args.if_not_exists() || args.if_match().is_some() {
346            return Err(opendal::Error::new(
347                ErrorKind::Unsupported,
348                "conditional copy is not supported by the HDFS fallback",
349            ));
350        }
351
352        let inner = Arc::clone(&self.inner);
353        let context = ctx.clone();
354        let renamer = self.renamer.clone();
355        let from = from.to_string();
356        let to = to.to_string();
357        Ok(oio::OneShotCopier::new(async move {
358            copy_via_read_write(inner, &context, renamer, &from, &to).await
359        }))
360    }
361
362    fn delete(&self, ctx: &OperationContext) -> Result<Self::Deleter> {
363        self.inner.delete(ctx)
364    }
365
366    fn list(&self, ctx: &OperationContext, path: &str, args: OpList) -> Result<Self::Lister> {
367        self.inner.list(ctx, path, args)
368    }
369
370    fn compose(&self, ctx: &OperationContext, to: &str, args: OpCompose) -> Result<Self::Composer> {
371        self.inner.compose(ctx, to, args)
372    }
373
374    async fn rename(
375        &self,
376        ctx: &OperationContext,
377        from: &str,
378        to: &str,
379        args: OpRename,
380    ) -> Result<RpRename> {
381        self.inner.rename(ctx, from, to, args).await
382    }
383
384    async fn presign(
385        &self,
386        ctx: &OperationContext,
387        path: &str,
388        args: OpPresign,
389    ) -> Result<RpPresign> {
390        self.inner.presign(ctx, path, args).await
391    }
392}
393
394async fn copy_via_read_write(
395    inner: Servicer,
396    context: &OperationContext,
397    renamer: AtomicRenamer,
398    source_path: &str,
399    target_path: &str,
400) -> Result<Metadata> {
401    let reader = inner.read(context, source_path, OpRead::new())?;
402    let (_, mut reader) = reader.open(opendal::BytesRange::from(..)).await?;
403    let temporary_path = temporary_path(target_path);
404    let mut writer = inner.write(context, &temporary_path, OpWrite::new())?;
405
406    loop {
407        let buffer = match reader.read().await {
408            Ok(buffer) => buffer,
409            Err(error) => {
410                abort_and_delete(&inner, context, writer, &temporary_path).await;
411                return Err(error);
412            }
413        };
414        if buffer.is_empty() {
415            break;
416        }
417        if let Err(error) = writer.write(buffer).await {
418            abort_and_delete(&inner, context, writer, &temporary_path).await;
419            return Err(error);
420        }
421    }
422
423    let metadata = match writer.close().await {
424        Ok(metadata) => metadata,
425        Err(error) => {
426            drop(writer);
427            let _ = delete_path(&inner, context, &temporary_path).await;
428            return Err(error);
429        }
430    };
431    drop(writer);
432
433    if let Err(error) = renamer
434        .rename(&inner, context, &temporary_path, target_path)
435        .await
436    {
437        let _ = delete_path(&inner, context, &temporary_path).await;
438        return Err(error);
439    }
440
441    Ok(metadata)
442}
443
444async fn abort_and_delete(
445    inner: &Servicer,
446    context: &OperationContext,
447    mut writer: oio::Writer,
448    path: &str,
449) {
450    if let Err(error) = writer.abort().await {
451        warn!(error; "Failed to abort HDFS writer for temporary file {path}");
452    }
453    drop(writer);
454    if let Err(error) = delete_path(inner, context, path).await {
455        warn!(error; "Failed to delete temporary file {path} after a failed HDFS write");
456    }
457}
458
459async fn delete_path(inner: &Servicer, context: &OperationContext, path: &str) -> Result<()> {
460    let mut deleter = inner.delete(context)?;
461    deleter.delete(path, OpDelete::new()).await?;
462    deleter.close().await
463}
464
465// TODO(fengjiachun): Clean up temporary files left behind after a process crash.
466fn temporary_path(path: &str) -> String {
467    let suffix = format!(".greptime-{}.tmp", Uuid::new_v4());
468    match path.rsplit_once('/') {
469        Some((parent, name)) => format!("{parent}/.{name}{suffix}"),
470        None => format!(".{path}{suffix}"),
471    }
472}
473
474#[cfg(test)]
475mod tests {
476    use opendal::services::Fs;
477    use opendal::{Operator, Writer};
478    use tempfile::TempDir;
479
480    use super::*;
481    use crate::layers::mock::{MockLayerBuilder, MockWriterFactory};
482
483    fn test_store() -> (TempDir, Operator) {
484        let directory = tempfile::tempdir().unwrap();
485        let store = Operator::new(Fs::default().root(directory.path().to_str().unwrap()))
486            .unwrap()
487            .layer(HdfsCompatibilityLayer::new_for_test());
488        (directory, store)
489    }
490
491    #[tokio::test]
492    async fn test_service_operations() {
493        let (_directory, store) = test_store();
494        store.create_dir("data/").await.unwrap();
495        store.write("data/source", "contents").await.unwrap();
496        assert_eq!(8, store.stat("data/source").await.unwrap().content_length());
497        assert_eq!(
498            b"onte",
499            store
500                .read_with("data/source")
501                .range(1..5)
502                .await
503                .unwrap()
504                .to_bytes()
505                .as_ref()
506        );
507        store.rename("data/source", "data/target").await.unwrap();
508        let entries = store.list("data/").await.unwrap();
509        assert!(entries.iter().any(|entry| entry.path() == "data/target"));
510        assert!(!store.exists("data/source").await.unwrap());
511        store.delete("data/target").await.unwrap();
512        assert!(!store.exists("data/target").await.unwrap());
513    }
514
515    #[tokio::test]
516    async fn test_atomic_write_keeps_old_data_after_abort() {
517        let (_directory, store) = test_store();
518        store.write("manifest.json", "old").await.unwrap();
519
520        let mut writer: Writer = store.writer("manifest.json").await.unwrap();
521        writer.write("new").await.unwrap();
522        writer.abort().await.unwrap();
523
524        assert_eq!(
525            b"old",
526            store
527                .read("manifest.json")
528                .await
529                .unwrap()
530                .to_bytes()
531                .as_ref()
532        );
533        assert!(
534            store
535                .list("")
536                .await
537                .unwrap()
538                .iter()
539                .all(|entry| !entry.path().contains(".greptime-"))
540        );
541    }
542
543    #[tokio::test]
544    async fn test_atomic_write_replaces_on_close() {
545        let (_directory, store) = test_store();
546        store.write("manifest.json", "old").await.unwrap();
547        store.write("manifest.json", "new").await.unwrap();
548
549        assert_eq!(
550            b"new",
551            store
552                .read("manifest.json")
553                .await
554                .unwrap()
555                .to_bytes()
556                .as_ref()
557        );
558        assert!(
559            store
560                .list("")
561                .await
562                .unwrap()
563                .iter()
564                .all(|entry| !entry.path().contains(".greptime-"))
565        );
566    }
567
568    #[tokio::test]
569    async fn test_copy_fallback_streams_to_target() {
570        let (_directory, store) = test_store();
571        store.write("source.parquet", "contents").await.unwrap();
572        store
573            .copy("source.parquet", "target.parquet")
574            .await
575            .unwrap();
576
577        assert_eq!(
578            b"contents",
579            store
580                .read("target.parquet")
581                .await
582                .unwrap()
583                .to_bytes()
584                .as_ref()
585        );
586    }
587
588    /// A writer that starts failing once more than `fail_after` bytes have been
589    /// accepted, to simulate a failure in the middle of an atomic write.
590    struct FailingWriter {
591        inner: oio::Writer,
592        written: usize,
593        fail_after: usize,
594    }
595
596    impl oio::Write for FailingWriter {
597        async fn write(&mut self, buffer: Buffer) -> Result<()> {
598            if self.written + buffer.len() > self.fail_after {
599                return Err(opendal::Error::new(
600                    ErrorKind::Unexpected,
601                    "injected write failure",
602                ));
603            }
604            self.written += buffer.len();
605            self.inner.write(buffer).await
606        }
607
608        async fn close(&mut self) -> Result<Metadata> {
609            self.inner.close().await
610        }
611
612        async fn abort(&mut self) -> Result<()> {
613            self.inner.abort().await
614        }
615    }
616
617    #[tokio::test]
618    async fn test_atomic_write_cleans_up_temporary_file_after_write_error() {
619        let directory = tempfile::tempdir().unwrap();
620        let factory: MockWriterFactory = Arc::new(|path, _args, inner| {
621            if path.contains(".greptime-") {
622                Box::new(FailingWriter {
623                    inner,
624                    written: 0,
625                    fail_after: 4,
626                })
627            } else {
628                inner
629            }
630        });
631        let mock_layer = MockLayerBuilder::default()
632            .writer_factory(factory)
633            .build()
634            .unwrap();
635        let store = Operator::new(Fs::default().root(directory.path().to_str().unwrap()))
636            .unwrap()
637            .layer(mock_layer)
638            .layer(HdfsCompatibilityLayer::new_for_test());
639
640        store.write("data/target", "old").await.unwrap();
641
642        let mut writer: Writer = store.writer("data/target").await.unwrap();
643        writer.write("part").await.unwrap();
644
645        let error = writer.write("ial").await.unwrap_err();
646        assert_eq!(ErrorKind::Unexpected, error.kind());
647        assert!(error.to_string().contains("injected write failure"));
648
649        // The destination keeps its previous contents...
650        assert_eq!(
651            b"old",
652            store.read("data/target").await.unwrap().to_bytes().as_ref()
653        );
654        // ...and the temporary file is removed even though the writer was not
655        // closed or aborted by the caller.
656        assert!(
657            store
658                .list("data/")
659                .await
660                .unwrap()
661                .iter()
662                .all(|entry| !entry.path().contains(".greptime-"))
663        );
664    }
665}
object_store/layers/hdfs.rs

object_store/layers/
hdfs.rs
