[sheepdog] [PATCH 3/3] zk_control: add remove method

[Liu Yuan]

With remove method, we can remove nasty 5s sleep in the _kill_sheep(), will
boost tests a lot with zk.

Signed-off-by: Liu Yuan <namei.unix at gmail.com>
---
 tests/functional/common.rc |   11 +++++--
 tools/zk_control.c         |   68 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/tests/functional/common.rc b/tests/functional/common.rc
index dfd47e6..b73f987 100644
--- a/tests/functional/common.rc
+++ b/tests/functional/common.rc
@@ -156,7 +156,10 @@ _cleanup()
     done
 
     if [[ "$DRIVER" == zoo* ]];then
-        sleep 5 # wait for zookeeper ephemeral node to be deleted
+        ../../tools/zk_control remove "/sheepdog/member"
+        if [ $? -ne 0 ]; then
+            _die "failed to remove members"
+        fi
     fi
 }
 
@@ -324,7 +327,11 @@ _kill_sheep()
     _wait_for_sheep_stop $1
 
     if [[ "$DRIVER" == zoo* ]];then
-	    _kill_zk_session $1
+        local path="/sheepdog/member/IPv4 ip:127.0.0.1 port:$((7000+$1))"
+        ../../tools/zk_control remove "$path"
+        if [ $? -ne 0 ]; then
+            _die "failed to remove $path"
+        fi
     fi
 }
 
diff --git a/tools/zk_control.c b/tools/zk_control.c
index 30050b5..6b3b136 100644
--- a/tools/zk_control.c
+++ b/tools/zk_control.c
@@ -14,9 +14,37 @@
 #include <zookeeper/zookeeper.h>
 #include <string.h>
 
+#define FOR_EACH_ZNODE(parent, path, strs)			       \
+	for ((strs)->data += (strs)->count;			       \
+	     (strs)->count-- ?					       \
+		     snprintf(path, sizeof(path), "%s/%s", parent,     \
+			      *--(strs)->data) : (free((strs)->data), 0); \
+	     free(*(strs)->data))
+
 static const char *hosts = "127.0.0.1:2181";
 static zhandle_t *zk_handle;
 
+static inline ZOOAPI int zk_delete_node(const char *path)
+{
+	int rc;
+	do {
+		rc = zoo_delete(zk_handle, path, -1);
+	} while (rc == ZOPERATIONTIMEOUT || rc == ZCONNECTIONLOSS);
+
+	return rc;
+}
+
+static inline ZOOAPI int zk_get_children(const char *path,
+					 struct String_vector *strings)
+{
+	int rc;
+	do {
+		rc = zoo_get_children(zk_handle, path, 1, strings);
+	} while (rc == ZOPERATIONTIMEOUT || rc == ZCONNECTIONLOSS);
+
+	return rc;
+}
+
 static int do_kill(int argc, char **argv)
 {
 	char *path;
@@ -59,7 +87,47 @@ static int do_kill(int argc, char **argv)
 
 static int do_remove(int argc, char **argv)
 {
+	struct String_vector strs;
+	int rc;
+	char *node;
+	char path[256];
+
+	if (argc != 3) {
+		fprintf(stderr, "remove: need specify path\n");
+		return -1;
+	}
+
+	node = argv[2];
+
+	rc = zk_get_children(node, &strs);
+	switch (rc) {
+	case ZOK:
+		FOR_EACH_ZNODE(node, path, &strs) {
+			rc = zk_delete_node(path);
+			if (rc != ZOK && rc != ZNONODE) {
+				fprintf(stderr, "failed to delete child "
+					"%s, %s\n",
+					path, zerror(rc));
+				goto err;
+			}
+		}
+		break;
+	case ZNOCHILDRENFOREPHEMERALS:
+		break;
+	case ZNONODE:
+		return 0;
+	default:
+		goto err;
+	}
+
+	rc = zk_delete_node(node);
+	if (rc != ZOK && rc != ZNONODE)
+		goto err;
+
 	return 0;
+err:
+	fprintf(stderr, "failed to delete %s, %s\n", node, zerror(rc));
+	return -1;
 }
 
 static struct control_handler {
-- 
1.7.9.5