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

use acl::*;
use consts::{CreateMode, ZkError};
use zookeeper::{ZkResult, ZooKeeper};
use std::iter::once;
use std::collections::VecDeque;

/// Extended ZooKeeper operations that are not needed for the "core."
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`.
    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.
    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.
    fn delete_recursive(&self, path: &str) -> ZkResult<()>;
}

impl ZooKeeperExt for ZooKeeper {
    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) {
                Ok(_) | Err(ZkError::NodeExists) => {}
                Err(e) => return Err(e),
            }
        }
        Ok(())
    }

    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)?;
            children
                .into_iter()
                .map(|child| format!("{}/{}", current, child))
                .for_each(|full_path| {
                    result.push(full_path.clone());
                    queue.push_back(full_path);
                });
        }

        Ok(result)
    }

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

        Ok(())
    }
}