import sys, socket, time, logging
import shlex, subprocess
from kazoo.client import KazooClient

def zkConnect(conn_str):
  zk = KazooClient(hosts=conn_str)
  zk.start()
  return zk

def zkCreateJobDir(zk, job_name):
  zk.ensure_path("/napper/memcached/%s" % (job_name))

def zkRemoveJobDir(zk, job_name):
  zk.delete("/napper/memcached/%s" % (job_name), recursive=True)

def zkRegisterWorker(zk, job_name, hostname, port):
  print "Registering myself as %s:%d" % (hostname, port)
  while zk.exists("/napper/memcached/%s:%d" % (hostname, port)):
    port += 1
  zk.create("/napper/memcached/%s:%d" % (hostname, port), "%d" % (port))
  zk.create("/napper/memcached/%s/%s:%d" % (job_name, hostname, port), "%d" % (port))
  return port

logging.basicConfig()

if len(sys.argv) < 6:
  print "usage: napper_naiad <Zookeeper hostname:port> <job name> <worker ID> <num workers> <executable>"
  sys.exit(1)

hostport = sys.argv[1]
job_name = sys.argv[2]
worker_id = int(sys.argv[3])
num_workers = int(sys.argv[4])
memcached_path = " ".join(sys.argv[5:])

client = zkConnect(hostport)
zkCreateJobDir(client, job_name)

actual_port = zkRegisterWorker(client, job_name, socket.gethostname(), 11211)

done = False
hosts = []
while not done:
  children = client.get_children("/napper/memcached/%s" % (job_name))
  if len(children) == num_workers:
    print "All workers are here!"
    for c in children:
      data, stat = client.get("/napper/memcached/%s/%s" % (job_name, c))
      print "%s:%s" % (c, data)
    done = True
  time.sleep(1)

# execute program
command = "%s -m 1024 -p %d" % (memcached_path, actual_port)
print "RUNNING: %s" % (command)
subprocess.call(shlex.split(command))

# clean up afterwards
if worker_id == 0:
  for c in children:
    data, stat = client.get("/napper/memcached/%s/%s" % (job_name, c))
    print "Unregistering %s:%s" % (c, data)
    client.delete("/napper/memcached/%s", c, recursive=True)
  zkRemoveJobDir(client, job_name)
  print "Deleted nodes for %s" % (job_name)
client.stop()

sys.exit(0)

@@ -88,16 +88,61 @@
 ooClient
+%0Afrom kazoo.exceptions import NodeExistsError
 %0A%0Adef zk
@@ -665,24 +665,40 @@
 %25d%22 %25 (port)
+, ephemeral=True
 )%0A  zk.creat
@@ -771,16 +771,32 @@
 %25 (port)
+, ephemeral=True
 )%0A  retu
@@ -1180,16 +1180,57 @@
 _name)%0A%0A
+done = False%0A%0Awhile not done:%0A  try:%0A    
 actual_p
@@ -1303,32 +1303,83 @@
 11)%0A
-%0A
+    
 done = 
-False%0Ahosts = %5B%5D
+True%0A  except NodeExistsError:%0A    pass%0A%0Ahosts = %5B%5D%0Adone = False
 %0Awhi
@@ -1511,22 +1511,20 @@
 nt %22All 
-worker
+task
 s are he
@@ -1845,322 +1845,48 @@
 %0A%0A# 
-clean up afterwards%0Aif worker_id == 0:%0A  for c in children:%0A    data, stat = client.get(%22/napper/memcached/%25s/%25s%22 %25 (job_name, c))%0A    print %22Unregistering %25s:%25s%22 %25 (c, data)%0A    client.delete(%22/napper/memcached/%25s%22, c, recursive=True)%0A  zkRemoveJobDir(client, job_name)%0A  print %22Deleted nodes for %25s%22 %25 (job_name)
+this will implicitly clean up afterwards
 %0Acli

Fix race condition in port allocation for memcached.