diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..c18dd8d --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +__pycache__/ diff --git a/barrier.py b/barrier.py index e399ef6..4d621b6 100755 --- a/barrier.py +++ b/barrier.py @@ -3,27 +3,49 @@ import argparse import socket -def barrier(host, port, wait_for): - with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: - s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - s.bind((host, port)) - s.listen() +class Barrier: + def __init__(self, host, port, wait_for): + self.host = host + self.port = port + self.wait_for = wait_for + def listen(self): + self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + self.sock.bind((self.host, self.port)) + self.sock.listen() + + def waitSingle(self): connections = [] addresses = [] while True: - conn, addr = s.accept() + conn, addr = self.sock.accept() connections.append(conn) addresses.append(addr) - print("Connection from {}".format(addr)) + yield addr - if len(connections) == wait_for: + if len(connections) == self.wait_for: break for conn in connections: conn.close() + return None + + def wait(self): + g = self.waitSingle() + + conn = [] + while True: + try: + conn.append(next(g)) + except StopIteration: + break + + return conn + + if __name__ == "__main__": parser = argparse.ArgumentParser() @@ -52,5 +74,16 @@ if __name__ == "__main__": results = parser.parse_args() + barrier = Barrier(results.host, results.port, results.processes) + barrier.listen() print("Barrier listens on {}:{} and waits for {} processes".format(results.host, results.port, results.processes)) - barrier(results.host, results.port, results.processes) \ No newline at end of file + + # connectedAddr = barrier.wait() + connectedAddrGen = barrier.waitSingle() + + while True: + try: + connectedAddr = next(connectedAddrGen) + print("Connection from {}".format(connectedAddr)) + except StopIteration: + break \ No newline at end of file