diff --git a/confluent_client/bin/nodeconsole b/confluent_client/bin/nodeconsole index 29e9f204..fe6decc9 100755 --- a/confluent_client/bin/nodeconsole +++ b/confluent_client/bin/nodeconsole @@ -910,7 +910,7 @@ async def do_screenshot(): dorefresh = False else: dorefresh = True - time.sleep(options.interval) + await asyncio.sleep(options.interval) sys.exit(0) async def grab_vncs(urlbynode): diff --git a/confluent_client/bin/nodediscover b/confluent_client/bin/nodediscover index f4e2691a..4c45a933 100755 --- a/confluent_client/bin/nodediscover +++ b/confluent_client/bin/nodediscover @@ -20,7 +20,6 @@ import csv import optparse import os import sys -import time path = os.path.dirname(os.path.realpath(__file__)) path = os.path.realpath(os.path.join(path, '..', 'lib', 'python')) @@ -51,10 +50,11 @@ columnmapping = { } #TODO: add chassis uuid -def register_endpoint(options, session, addr): +async def register_endpoint(options, session, addr): neednewline = False current = 0 - for rsp in session.update('/discovery/register', {'addresses': addr}): + total = 0 + async for rsp in session.update('/discovery/register', {'addresses': addr}): if 'count' in rsp: total = rsp['count'] elif total > 1: @@ -70,15 +70,15 @@ def register_endpoint(options, session, addr): if neednewline: print('') -def subscribe_discovery(options, session, subscribe, targ): +async def subscribe_discovery(options, session, subscribe, targ): keyn = 'subscribe' if subscribe else 'unsubscribe' payload = {keyn: targ} if subscribe: - for rsp in session.update('/discovery/subscriptions/{0}'.format(targ), payload): + async for rsp in session.update('/discovery/subscriptions/{0}'.format(targ), payload): if 'status' in rsp: print(rsp['status']) else: - for rsp in session.delete('/discovery/subscriptions/{0}'.format(targ)): + async for rsp in session.delete('/discovery/subscriptions/{0}'.format(targ)): if 'status' in rsp: print(rsp['status']) @@ -180,6 +180,9 @@ def datum_to_attrib(datum): unique_fields = frozenset(['serial', 'mac', 'uuid']) +# Cap how many nodes hold a discovery session at once while importing +maxconcurrentassign = 128 + async def import_csv(options, session): nodedata = [] unique_data = {} @@ -211,12 +214,12 @@ async def import_csv(options, session): alldata.append(nodedatum) allthere = True for nodedatum in alldata: - if not search_record(nodedatum, options, session) and not broken: + if not await search_record(nodedatum, options, session) and not broken: allthere = False await blocking_scan(session) break for nodedatum in alldata: - if not allthere and not search_record(nodedatum, options, session): + if not allthere and not await search_record(nodedatum, options, session): sys.stderr.write( "Could not match the following data: " + repr(nodedatum) + '\n') @@ -224,11 +227,13 @@ async def import_csv(options, session): nodedata.append(nodedatum) if broken: sys.exit(1) + assignments = [] + assignlimit = asyncio.Semaphore(maxconcurrentassign) for datum in nodedata: - maclist = search_record(datum, options, session) + maclist = await search_record(datum, options, session) datum = datum_to_attrib(datum) nodename = datum['name'] - for res in session.create('/nodes/', datum): + async for res in session.create('/nodes/', datum): if 'error' in res: sys.stderr.write(res['error'] + '\n') exitcode |= res.get('errorcode', 1) @@ -237,13 +242,26 @@ async def import_csv(options, session): print('Defined ' + res['created']) else: print(repr(res)) - child = os.fork() - if child: - continue + assignments.append( + asyncio.create_task(assign_macs(maclist, nodename, assignlimit))) + for rcode in await asyncio.gather(*assignments, return_exceptions=True): + if isinstance(rcode, BaseException): + sys.stderr.write('Error assigning discovery data: {0}\n'.format(rcode)) + rcode = 1 + exitcode |= rcode + if exitcode: + sys.exit(exitcode) + + +async def assign_macs(maclist, nodename, assignlimit): + exitcode = 0 + async with assignlimit: + # A session of our own, since the connection carries one request at a + # time and the caller's is busy defining the remaining nodes + mysess = client.Command() for mac in maclist: - mysess = client.Command() - for res in mysess.update('/discovery/by-mac/{0}'.format(mac), - {'node': nodename}): + async for res in mysess.update('/discovery/by-mac/{0}'.format(mac), + {'node': nodename}): if 'error' in res: sys.stderr.write(res['error'] + '\n') exitcode |= res.get('errorcode', 1) @@ -252,14 +270,7 @@ async def import_csv(options, session): print('Discovered ' + res['assigned']) else: print(repr(res)) - sys.exit(0) - while True: - try: - os.wait() - except ChildProcessError: - break - if exitcode: - sys.exit(exitcode) + return exitcode async def list_discovery(options, session): @@ -364,7 +375,7 @@ async def assign_discovery(options, session, needid=True): async def blocking_scan(session): list([x async for x in session.update('/discovery/rescan', {'rescan': 'start'})]) while(list([x async for x in session.read('/discovery/rescan')])[0].get('scanning', False)): - time.sleep(0.5) + await asyncio.sleep(0.5) list([x async for x in session.update('/networking/macs/rescan', {'rescan': 'start'})]) @@ -427,11 +438,11 @@ async def main(): if args[0] == 'reassign': await assign_discovery(options, session, False) if args[0] == 'register': - register_endpoint(options, session, args[1]) + await register_endpoint(options, session, args[1]) if args[0] == 'subscribe': - subscribe_discovery(options, session, True, args[1]) + await subscribe_discovery(options, session, True, args[1]) if args[0] == 'unsubscribe': - subscribe_discovery(options, session, False, args[1]) + await subscribe_discovery(options, session, False, args[1]) if args[0] == 'rescan': await blocking_scan(session) print("Rescan complete") diff --git a/confluent_client/confluent/asynclient.py b/confluent_client/confluent/asynclient.py index 79b4308f..2f29e2e6 100644 --- a/confluent_client/confluent/asynclient.py +++ b/confluent_client/confluent/asynclient.py @@ -347,12 +347,12 @@ class Command(object): else: ikey = key if input is None: - for res in await self.read('/nodegroups/{0}/{1}'.format( + async for res in self.read('/nodegroups/{0}/{1}'.format( noderange, resource)): rc = self.handle_results(ikey, rc, res) else: kwargs[ikey] = input - for res in await self.update('/nodegroups/{0}/{1}'.format( + async for res in self.update('/nodegroups/{0}/{1}'.format( noderange, resource), kwargs): rc = self.handle_results(ikey, rc, res) return rc @@ -396,6 +396,7 @@ class Command(object): self.connection = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM) self.connection.setsockopt(socket.SOL_SOCKET, SO_PASSCRED, 1) self.connection.connect(self.serverloc) + self.connection.setblocking(False) async def _connect_tls(self): server, port = _parseserver(self.serverloc) @@ -667,10 +668,10 @@ def show_attr(attr, requestargs, seenattributes, options, node): return processattr -def printgroupattributes(session, requestargs, showtype, nodetype, noderange, options): +async def printgroupattributes(session, requestargs, showtype, nodetype, noderange, options): exitcode = 0 seenattributes = set([]) - for res in session.read('/{0}/{1}/attributes/{2}'.format(nodetype, noderange, showtype)): + async for res in session.read('/{0}/{1}/attributes/{2}'.format(nodetype, noderange, showtype)): if 'error' in res: sys.stderr.write(res['error'] + '\n') exitcode = 1 diff --git a/confluent_client/confluent/asynctlvdata.py b/confluent_client/confluent/asynctlvdata.py index 809880b5..e27b20d4 100644 --- a/confluent_client/confluent/asynctlvdata.py +++ b/confluent_client/confluent/asynctlvdata.py @@ -99,9 +99,9 @@ class ClientFile(object): -def _sendmsg(loop, fut, sock, msg, fds, rfd): - if rfd is not None: - loop.remove_reader(rfd) +def _sendmsg(loop, fut, sock, msg, fds, wfd): + if wfd is not None: + loop.remove_writer(wfd) if fut.cancelled(): return try: @@ -110,7 +110,7 @@ def _sendmsg(loop, fut, sock, msg, fds, rfd): [(socket.SOL_SOCKET, socket.SCM_RIGHTS, array.array("i", fds))]) except (BlockingIOError, InterruptedError): fd = sock.fileno() - loop.add_reader(fd, _sendmsg, loop, fut, sock, fd) + loop.add_writer(fd, _sendmsg, loop, fut, sock, msg, fds, fd) except Exception as exc: fut.set_exception(exc) else: @@ -127,6 +127,8 @@ def send_fds(sock, msg, fds): def _recvmsg(loop, fut, sock, msglen, maxfds, rfd): if rfd is not None: loop.remove_reader(rfd) + if fut.cancelled(): + return fds = array.array("i") # Array of ints try: msg, ancdata, flags, addr = sock.recvmsg(