1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81

use async_trait::async_trait;
use log::*;
use std::collections::VecDeque;
use std::iter::once;

use crate::{Acl, CreateMode, ZkError, ZkResult, ZooKeeper};

/// Extended ZooKeeper operations that are not needed for the "core."
#[async_trait]
pub trait ZooKeeperExt {
    /// Ensure that `path` exists and create all potential paths leading up to it if it does not.
    /// This operates in a manner similar to `mkdir -p`.
    async fn ensure_path(&self, path: &str) -> ZkResult<()>;

    /// Performs a breadth-first tree traversal of the tree starting at `path`,
    /// returning a list of fully prefixed child nodes.
    /// *NOTE*: This is not an atomic operation.
    async fn get_children_recursive(&self, path: &str) -> ZkResult<Vec<String>>;

    /// Deletes the node at `path` and all its children.
    /// *NOTE*: This is not an atomic operation.
    async fn delete_recursive(&self, path: &str) -> ZkResult<()>;
}

#[async_trait]
impl ZooKeeperExt for ZooKeeper {
    async fn ensure_path(&self, path: &str) -> ZkResult<()> {
        trace!("ensure_path {}", path);
        for (i, _) in path
            .chars()
            .chain(once('/'))
            .enumerate()
            .skip(1)
            .filter(|c| c.1 == '/')
        {
            match self
                .create(
                    &path[..i],
                    vec![],
                    Acl::open_unsafe().clone(),
                    CreateMode::Persistent,
                )
                .await
            {
                Ok(_) | Err(ZkError::NodeExists) => {}
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }

    async fn get_children_recursive(&self, path: &str) -> ZkResult<Vec<String>> {
        trace!("get_children_recursive {}", path);
        let mut queue: VecDeque<String> = VecDeque::new();
        let mut result = vec![path.to_string()];
        queue.push_front(path.to_string());

        while let Some(current) = queue.pop_front() {
            let children = self.get_children(&current, false).await?;
            children
                .into_iter()
                .map(|child| format!("{}/{}", current, child))
                .for_each(|full_path| {
                    result.push(full_path.clone());
                    queue.push_back(full_path);
                });
        }

        Ok(result)
    }

    async fn delete_recursive(&self, path: &str) -> ZkResult<()> {
        trace!("delete_recursive {}", path);
        let children = self.get_children_recursive(path).await?;
        for child in children.iter().rev() {
            self.delete(child, None).await?;
        }

        Ok(())
    }
}