Skip to main content

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