"""
This module provides a client class to connect to a Preseem GRPC service and
execute operations on it.  GRPC specific aspects are abstracted here, and
Python methods are provided.  By design, this module does not look at any
of the model objects, just the request/response objects and the rpc calls.

A large part of the responsibility of this module is to convert the third-party
GRPC library to the asyncio concurrency model.  This includes support for
streaming subscriptions, where the caller is blocked in a Subscribe call and
then has streaming messages asynchronously delivered into the event loop via
a callback.

This is designed to be easily replaceable by a stub client class for unit
testing.

All Preseem services are modeled here; I'm not super happy with that aspect
of the design, but for now it's left that way.
"""
import asyncio
import concurrent.futures
import functools
from importlib import import_module
import logging
import threading
import time

import grpc

from preseem.handoff_queue import HandoffQueue
from preseem_grpc_model import metrics_pb2, metrics_pb2_grpc
from preseem_grpc_model import network_metadata_pb2, network_metadata_pb2_grpc
from preseem_protobuf.config import company_configs_pb2, company_configs_pb2_grpc

# The reader thread enqueues decoded ops here and a single applier coroutine
# batch-drains them on the event loop, so the per-op cross-thread round-trip is
# amortized over a whole batch.  The queue is bounded so a slow applier exerts
# real backpressure on the reader (it stops draining the HTTP/2 receive window)
# instead of letting ops pile up unbounded in the client.  50k holds a large
# backlog comfortably; correctness does not depend on the queue staying small.
_SUBSCRIPTION_QUEUE_MAXSIZE = 50_000
# Cap ops applied per event-loop turn.  The steady SET apply path has no real
# await, so a whole batch runs in one event-loop turn; without a cap, draining a
# deep backlog (up to _SUBSCRIPTION_QUEUE_MAXSIZE) in one turn monopolizes the
# loop (~115ms for a 50k drain) and starves every other coroutine on it (other
# subscription appliers, periodic work).  Capping forces the applier to yield to
# the loop between batches.
#
# 1024 is the measured throughput/per-turn-time knee (60k-op soak):
#   cap   ops/s    per-turn   handoffs(60k ops)
#   256   ~289k    ~0.9ms     235
#   512   ~425k    ~1.2ms     118
#   1024  ~497k    ~2.0ms      59   <- knee
#   2048  ~515k    ~4.0ms      30
#   4096  ~552k    ~7.4ms      15
# Throughput keeps climbing past 512 because the per-hand-off run_in_executor
# round-trip is not yet negligible there; it flattens around 1024.  1024 keeps
# each loop turn ~2ms (small vs the coarse cadence of other loop work) while
# capturing ~497k ops/s, within ~10% of larger caps that cost 2-4x the turn time.
_MAX_DRAIN_BATCH = 1024

_APPLY_ERROR_LOG_INTERVAL = 10.0


class ConfigType():

    def __init__(self, name, prefix, value, pb, grpc):
        self.name = name  # the lowercase name of the enum
        self.value = value  # the int value of the enum
        self.prefix = prefix  # op prefix
        self.pb = pb  # pb2 import
        self.grpc = grpc  # pb2_grpc import
        self.svc = None  # service stub
        self.pb_type = None  # grpc type
        self.del_req = None  # delete request
        self.get_req = None  # get request
        self.set_req = None  # set request
        self.sub_req = None  # subscribe request


class Subscription(object):
    """Helper class to manage a GRPC subscription."""

    def __init__(self, client, op, req, cbk, handler=None, kind='unknown'):
        """
        Manage a GRPC subscription in asyncio.  Pass in the client object
        and the operation to call that will return the subscription iterator.
        If handler is provided, it is a generator that consumes the raw gRPC
        iterator and yields (op, data) tuples dispatched to cbk.  Otherwise,
        each message is passed directly to cbk(msg, syn).

        `kind` is a low-cardinality label identifying the subscription (e.g.
        'nm_ent', 'nm_ref', 'cfg:<type>') used to attribute log messages.
        """
        self._client = client
        self._op = op
        self._req = req
        self._cbk = cbk
        self._handler = handler
        self._kind = kind
        # A subscription-owned executor runs BOTH blocking jobs -- the reader
        # thread (self._run) and the blocking queue drain (self._drain_batch).
        # It must NOT share the process-global default executor: _drain_batch
        # parks in the hand-off queue's drain() until an item arrives, so routing
        # it through the default executor would tie up its bounded, shared thread
        # pool (also used by loop.getaddrinfo() for DNS) across every long-lived
        # subscription, starving DNS and other subscriptions' appliers.
        # max_workers=2: one thread for the reader, one for the drain.
        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=2)
        self._loop = asyncio.get_event_loop()
        self._iter = None
        # Bounded, cancellation-aware hand-off from the reader thread to the
        # applier coroutine.  Each item is the args tuple to apply via self._cbk.
        # Its blocking put/drain wake out-of-band on stop()/producer_done(), so
        # neither worker is ever parked past shutdown -- no timeout poll, no
        # in-band sentinel.
        self._queue = HandoffQueue(_SUBSCRIPTION_QUEUE_MAXSIZE)
        # Set by close()/teardown to stop the reader's NON-queue waits (the
        # reconnect/create _backoff and the loop guards).  The queue's own stop()
        # handles a reader parked in put() and a drain parked in drain().
        self._closing = threading.Event()
        # Throttle state for callback-apply error logging (see
        # _record_apply_error).  Touched only on the applier coroutine.  None
        # means "no error logged yet".
        self._apply_error_last_log = None
        self._apply_errors_since_log = 0

    def _request_stop(self):
        """Force both workers to exit promptly, from any thread, idempotently.

        Every blocking point the workers can be parked in observes the stop:
          - the reader in the queue's put() on a full queue, and the applier's
            drain parked in the queue's drain() -> self._queue.stop() wakes both
            out-of-band (instantly), with no timeout poll;
          - the reader in a reconnect/error backoff -> the closing flag, which
            _backoff waits on (returning immediately when set) vs a blind sleep;
          - the reader creating the subscription -> a SINGLE attempt (no internal
            retry sleep close() could not observe), so a failure returns to the
            _closing-aware _backoff promptly;
          - the reader in the blocking gRPC read (self._iter.next() / inside the
            handler) -> iterator cancel.  Cancel only works if self._iter is
            already set; if close() lands while it is still None (just before the
            read loop), the re-check guarding the read in _run catches the flag.
        We never enqueue anything to signal stop: termination is out-of-band, so
        there is no in-band item a concurrent drainer could swallow."""
        self._closing.set()
        self._queue.stop()
        it = self._iter  # snapshot: the reader may rebind _iter across reconnects
        if it is not None:
            try:
                it.cancel()
            except Exception:
                pass

    def close(self):
        """Close the subscription; both worker threads terminate immediately --
        the cancellation-aware hand-off queue wakes a parked put/drain at once --
        even under backpressure or while idle."""
        self._request_stop()

    def _put(self, item):
        """Enqueue from the reader thread.  Blocks while the queue is full to
        exert backpressure, but the queue wakes out-of-band on stop(), so close()
        unblocks a parked reader instantly.  Returns False if the subscription is
        closing and the item was dropped."""
        return self._queue.put(item)

    def _backoff(self, seconds):
        """Interruptible reader backoff.  Waits up to `seconds`, but returns
        immediately once close() sets the closing flag, so a reader sleeping
        between reconnect attempts wakes promptly on shutdown instead of after
        the full delay.  Returns False if the subscription is closing."""
        # Event.wait returns True iff the flag is set, i.e. we should stop.
        return not self._closing.wait(seconds)

    def _run(self):
        """Read from the subscription, in a separate thread.

        The reader only DECODES and ENQUEUES ops; it never blocks on
        a per-op cross-thread round-trip, so it keeps draining the HTTP/2
        receive window as fast as the server sends.  The bounded queue provides
        backpressure if the applier falls behind.  On exit it marks the hand-off
        queue producer_done so the applier terminates out-of-band (no in-band
        sentinel).  A reader parked in the queue's put() on a full queue is woken
        by close() (stop())."""
        retry = False
        try:
            while not self._closing.is_set():
                # Setup the subscription
                syn = False
                if self._iter is None:
                    try:
                        if retry:
                            logging.info('retry subscription to: %s', self._client.host)
                        else:
                            retry = True
                        # ONE attempt only -- no internal retry loop here.  A
                        # retrying create would sleep uninterruptibly (close()
                        # can neither cancel nor be observed while _iter is None)
                        # and could return None on exhaustion (-> None.next()
                        # deref).  Letting a single attempt raise routes retries
                        # through _run's own _closing-aware _backoff below.
                        self._iter = self._op(self._req)
                    except Exception as e:
                        logging.warning('Error creating subscription: %s', e)
                        if not self._backoff(10):
                            return  # closing; stop reading
                        continue
                    syn = True
                    # A successful create can still return None (defensive);
                    # never fall through into None.next().
                    if self._iter is None:
                        if not self._backoff(10):
                            return  # closing; stop reading
                        continue
                    # Close the TOCTOU window: a close() that ran while _iter was
                    # still None could not cancel it, so re-check here before we
                    # block in the read.  Any close() AFTER this assignment sees
                    # _iter and cancels it, interrupting the read.
                    if self._closing.is_set():
                        return

                try:
                    if self._handler:
                        for op in self._handler(self._iter):
                            if not self._put(op):
                                return  # closing; stop reading
                    else:
                        msg = self._iter.next()
                except grpc.RpcError as e:
                    # https://grpc.io/grpc/python/grpc.html?highlight=method#grpc.StatusCode
                    # Disconnect, need to restart subscription.  Restart
                    # the initializing state after this.
                    if e.code() == grpc.StatusCode.CANCELLED and e.details(
                    ) == 'Locally cancelled by application!':
                        return  # cancelled by application
                    elif e.code() == grpc.StatusCode.NOT_FOUND:
                        # STM-4023 the thing we are subscribed to has gone away.
                        logging.warning('GRPC NotFound error: %s, exit subscription',
                                        e.details())
                        break
                    logging.warning('GRPC ERROR: %s: %s for host: %s', e.code(),
                                    e.details(), self._client.host)
                    logging.info('Reconnect to %s after 10s', self._client.host)
                    self._iter = None
                    if not self._backoff(10):
                        return  # closing; stop reading
                    continue

                except Exception as e:
                    logging.exception('Error reading from subscription '
                                      'iterator: %s', e)
                    self._iter = None  # drop it so the next loop reconnects
                    if not self._backoff(10):
                        return  # closing; stop reading
                    continue
                if not self._handler:
                    if not self._put((msg, syn)):
                        return  # closing; stop reading
        finally:
            # The reader has produced everything it will; tell the hand-off queue
            # so a drain blocked on an empty queue wakes and reports done once the
            # backlog is applied.  Out-of-band: no in-band sentinel to swallow.
            self._queue.producer_done()

    def _drain_batch(self):
        """Block until items are available (or the queue is stopped / the reader
        is done), then take up to _MAX_DRAIN_BATCH ops in FIFO order.  Returns
        (batch, done) from the hand-off queue: done is True once stopped, or once
        the reader finished and the queue is drained.  The cap bounds per-event-
        loop-turn work; the applier yields between batches and is re-invoked until
        done.  The queue wakes this immediately on a put and out-of-band on
        stop()/producer_done(), so there is no timeout and no idle poll."""
        return self._queue.drain(_MAX_DRAIN_BATCH)

    def _record_apply_error(self):
        """Record a callback-apply failure (op skipped) without spamming the log.

        Logs the full traceback only on the first failure and then at most once
        per _APPLY_ERROR_LOG_INTERVAL, folding the suppressed count into the
        throttled summary.  This avoids the per-op traceback firehose a recurring
        failure (e.g. a poison op) would otherwise produce at batch-drain rate."""
        self._apply_errors_since_log += 1
        now = time.monotonic()
        first = self._apply_error_last_log is None
        if not first and now - self._apply_error_last_log < _APPLY_ERROR_LOG_INTERVAL:
            return
        count = self._apply_errors_since_log
        self._apply_error_last_log = now
        self._apply_errors_since_log = 0
        if first:
            logging.exception(
                'Error applying subscription callback (op skipped); '
                'further errors throttled to once per %ss', _APPLY_ERROR_LOG_INTERVAL)
        else:
            logging.exception(
                'Error applying subscription callback (op skipped); '
                '%d errors in the last ~%ss', count, _APPLY_ERROR_LOG_INTERVAL)

    def _flush_apply_errors(self):
        """Surface any apply errors counted since the last throttled summary.

        The throttle only emits a summary on the NEXT error after the interval,
        so a transient burst that then stops would leave its tail sitting in
        _apply_errors_since_log, unreported in the log.  Called on applier
        teardown to flush that tail.  No traceback: there is no active exception
        at teardown."""
        if self._apply_errors_since_log:
            count = self._apply_errors_since_log
            self._apply_errors_since_log = 0
            logging.error(
                'Error applying subscription callback (op skipped); '
                '%d further error(s) before the subscription stopped', count)

    async def _apply(self):
        """Drain the queue in batches and apply each op with a plain await,
        amortizing the cross-thread round-trip over the whole batch.  Ops are
        applied in enqueue order; a failing callback is logged and skipped.
        Batches are capped (_MAX_DRAIN_BATCH) so the event loop gets a yield
        point between them.  Terminates out-of-band once the reader is done and
        the queue is drained, or immediately once close()/stop() fires.

        The hand-off queue's drain blocks until an item is available and wakes
        immediately on a put (so delivery latency is untouched) and out-of-band
        on stop()/producer_done() (so there is no idle poll and teardown is
        instant -- a parked drain is released the moment stop() runs)."""
        try:
            while not self._closing.is_set():
                batch, done = await self._loop.run_in_executor(
                    self._executor, self._drain_batch)
                for op in batch:
                    try:
                        await self._cbk(*op)
                    except asyncio.CancelledError:
                        # On Python 3.7 CancelledError subclasses Exception, so
                        # it would otherwise be eaten by the broad except below
                        # and hang shutdown.  Propagate cancellation.
                        raise
                    except Exception:
                        self._record_apply_error()
                if done:
                    return
        finally:
            # On cancellation (loop teardown) force the reader to stop so it
            # can't stay parked in _put on a full queue NOR in the blocking gRPC
            # read.  No queue draining -- termination is out-of-band.
            self._request_stop()
            # Flush any apply errors still pending in the throttle window so a
            # transient burst that stopped is not lost from the log.
            self._flush_apply_errors()

    async def run(self):
        """Read from the subscription and block the caller on the eventloop.

        Runs the reader (worker thread) and the batch-draining applier
        concurrently; the applier finishes once the reader has exited (marking
        the hand-off queue producer_done) and the queue is drained.

        If EITHER side raises, the closing flag is signalled so the other stops
        promptly, and both are awaited (return_exceptions=True) so neither future
        is orphaned or left as a "never retrieved" exception.  A CancelledError
        (loop teardown / caller cancellation) is re-raised so cancellation
        propagates cleanly.  Any OTHER unexpected exception is logged loudly with
        the subscription identity and NOT re-raised: the subscription stops, but
        the agent stays up rather than crashing or silently swallowing it."""
        reader = self._loop.run_in_executor(self._executor, self._run)
        applier = asyncio.ensure_future(self._apply())
        try:
            results = await asyncio.gather(reader, applier, return_exceptions=True)
        finally:
            # Whatever happened, force both workers to stop so neither lingers
            # parked in the hand-off queue nor in the blocking gRPC read.
            self._request_stop()
            # Both worker jobs are done (gather awaited them), so release the
            # executor's threads now rather than waiting for GC -- matters for
            # short-lived config subscriptions created over the client's life.
            self._executor.shutdown(wait=False)
        for result in results:
            if isinstance(result, asyncio.CancelledError):
                raise result
            if isinstance(result, BaseException):
                logging.exception(
                    'Unexpected error in %s subscription to %s; subscription '
                    'stopped',
                    self._kind,
                    self._client.host,
                    exc_info=result)


class PreseemGrpcClient(object):
    """Implement the GRPC client to the Preseem API, adapted to asyncio."""

    def __init__(self, api_key, host=None, port=None, compression=False, insecure=None):
        self.api_key = api_key
        self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=10)
        self.host = host or 'nmapi.preseem.com'
        self.port = port or 443
        self.compression = compression
        self.insecure = insecure
        self._channel = None
        self._metrics_svc = None
        self._nm_svc = None
        self._nm_ref_svc = None
        self._config2_svc = None

        self._nm_ent_subscription = None
        self._nm_ref_subscription = None
        self._cfg_subscriptions = {}
        self._cfg_lock = None

        self.config_types = {}
        for ctname, ctid in company_configs_pb2.Type.items():
            if ctid < 2:
                continue  # ignore NODE configs for now
            try:
                cn = ctname.lower()
                pmod = import_module(f'preseem_protobuf.config.{cn}_pb2')
                gmod = import_module(f'preseem_protobuf.config.{cn}_pb2_grpc')
                ct = self.config_types[ctid] = ConfigType(cn, cn.replace('_', ''), ctid,
                                                          pmod, gmod)
                # assume naming conventions for the operations
                capname = ''.join([x[:1].upper() + x[1:] for x in cn.split('_')])
                ct.pb_type = getattr(pmod, f'{capname}')
                ct.get_req = getattr(pmod, f'{capname}GetRequest')
                ct.set_req = getattr(pmod, f'{capname}SetRequest')
                ct.del_req = getattr(pmod, f'{capname}DeleteRequest')
                ct.sub_req = getattr(pmod, f'{capname}SubscribeRequest')
            except ModuleNotFoundError:
                # we don't support this, that's ok
                pass

    def close_subscriptions(self):
        if self._nm_ent_subscription:
            self._nm_ent_subscription.close()
        if self._nm_ref_subscription:
            self._nm_ref_subscription.close()
        for sub in self._cfg_subscriptions.values():
            sub.close()
        self._cfg_subscriptions.clear()

    def _grpc_op_collect(self, op, req, metadata, handler=None):
        """
        Run a GRPC operation and collect the results in this thread.
        If a handler function is passed in, that function will get an interator
        and can process it however it wants; whatever that function returns will be
        returned to the caller.
        """
        r = op(req, metadata=metadata)
        try:
            iterator = iter(r)  # this is to detect whether its a sequence
            if handler:
                return handler(iterator)
            return [x for x in iterator]  # this does network i/o
        except TypeError:  # not an iterable result, ok
            pass
        return r

    async def _grpc_op(self,
                       op,
                       req,
                       num_retries=5,
                       timeout=None,
                       collect_results=None,
                       metadata=None,
                       handler=None):
        loop = asyncio.get_event_loop()
        for retry_num in range(num_retries + 1):
            if retry_num > 0:
                await asyncio.sleep(2**retry_num)
                logging.debug('Retrying after error result')
            try:
                if timeout:
                    try:
                        if collect_results:
                            return await asyncio.wait_for(loop.run_in_executor(
                                self._executor, self._grpc_op_collect, op, req,
                                metadata, handler),
                                                          timeout=timeout)
                        else:
                            return await asyncio.wait_for(loop.run_in_executor(
                                self._executor,
                                functools.partial(op, req, metadata=metadata)),
                                                          timeout=timeout)
                    except asyncio.TimeoutError:
                        logging.warning('GRPC operation timed out: %s', op)
                        raise
                else:
                    if collect_results:
                        x = await loop.run_in_executor(self._executor,
                                                       self._grpc_op_collect, op, req,
                                                       metadata, handler)
                    else:
                        x = await loop.run_in_executor(
                            self._executor, functools.partial(op,
                                                              req,
                                                              metadata=metadata))
                    return x
            except grpc.RpcError as err:
                logging.debug('Error calling service: %s (%s)', err.code(),
                              err.details())
                if err.code() in (grpc.StatusCode.INTERNAL, grpc.StatusCode.NOT_FOUND,
                                  grpc.StatusCode.UNKNOWN,
                                  grpc.StatusCode.UNIMPLEMENTED,
                                  grpc.StatusCode.INVALID_ARGUMENT,
                                  grpc.StatusCode.PERMISSION_DENIED,
                                  grpc.StatusCode.RESOURCE_EXHAUSTED):
                    raise
                if retry_num == num_retries:
                    logging.warning('Error calling service: %s (%s)', err.code(),
                                    err.details())
                    raise
            except asyncio.CancelledError:
                # On Python 3.7 CancelledError subclasses Exception; never retry
                # or swallow cancellation -- propagate it so shutdown is prompt.
                raise
            except Exception as e:
                logging.warning('%s error calling service: %s', type(e).__name__, e)
                if retry_num == num_retries:
                    logging.exception('Error calling grpc service')
                    raise

    def _connect(self):
        """Connect to the server and authenticate. (blocking)"""

        class ApiKeyCallCredentials(grpc.AuthMetadataPlugin):

            def __init__(self, api_key):
                self.api_key = api_key

            def __call__(self, context, callback):
                callback((('api_key', self.api_key), ), None)

        class AuthInterceptor(
                grpc.UnaryUnaryClientInterceptor,
                grpc.UnaryStreamClientInterceptor,
                grpc.StreamUnaryClientInterceptor,
                grpc.StreamStreamClientInterceptor,
        ):
            """Universal interceptor that adds metadata to all RPC types."""

            def __init__(self, plugin):
                self.plugin = plugin

            def _augment_metadata(self, client_call_details):
                """Helper: calls your plugin and merges metadata."""
                md = []

                def callback(metadata, error):
                    if metadata:
                        md.extend(metadata)

                # Run the plugin to populate md
                self.plugin(None, callback)

                # Merge with any existing metadata
                combined = list(client_call_details.metadata or []) + md
                return client_call_details._replace(metadata=tuple(combined))

            # Unary → Unary
            def intercept_unary_unary(self, continuation, client_call_details, request):
                new_details = self._augment_metadata(client_call_details)
                return continuation(new_details, request)

            # Unary → Stream
            def intercept_unary_stream(self, continuation, client_call_details,
                                       request):
                new_details = self._augment_metadata(client_call_details)
                return continuation(new_details, request)

            # Stream → Unary
            def intercept_stream_unary(self, continuation, client_call_details,
                                       request_iterator):
                new_details = self._augment_metadata(client_call_details)
                return continuation(new_details, request_iterator)

            # Stream → Stream
            def intercept_stream_stream(self, continuation, client_call_details,
                                        request_iterator):
                new_details = self._augment_metadata(client_call_details)
                return continuation(new_details, request_iterator)

        dst = '{}:{}'.format(self.host, self.port)
        while self._nm_svc is None:
            try:
                # STM-3013, STM-3115 make grpc connection robust
                grpc_opts = [
                    ('grpc.keepalive_time_ms', 30000),  # send keepalive every 30s
                    ('grpc.keepalive_timeout_ms',
                     15000),  # keepalive ping timeout after 15s
                    ('grpc.keepalive_permit_without_calls',
                     True),  # send keepalives when no grpc calls
                    ('grpc.http2.max_pings_without_data', 0),
                    ('grpc.max_receive_message_length', 10 * 1024 * 1024),
                ]  # allow unlimited pings without data
                if self.compression:
                    grpc_opts.append(
                        ('grpc.default_compression_algorithm', grpc.Compression.Gzip))
                if self.api_key and not self.insecure:
                    call_creds = grpc.metadata_call_credentials(
                        ApiKeyCallCredentials(self.api_key))
                    try:
                        with open('/etc/pki/tls/certs/ca-bundle.crt', 'rb') as f:
                            ca_certs = f.read()
                    except:  # not fedora?  use default built-in CA bundle
                        ca_certs = None
                    ssl_creds = grpc.ssl_channel_credentials(root_certificates=ca_certs)
                    channel_creds = grpc.composite_channel_credentials(
                        ssl_creds, call_creds)
                    channel = grpc.secure_channel(dst, channel_creds, options=grpc_opts)
                else:
                    channel = grpc.insecure_channel(dst, options=grpc_opts)
                    if self.api_key:
                        channel = grpc.intercept_channel(
                            channel,
                            AuthInterceptor(ApiKeyCallCredentials(self.api_key)))
                self._channel = channel
                self._metrics_svc = metrics_pb2_grpc.NetworkMetricsStub(channel)
                self._nm_svc = network_metadata_pb2_grpc.NetworkMetadataServiceStub(
                    channel)
                self._nm_ref_svc = network_metadata_pb2_grpc.NetworkMetadataReferenceServiceStub(
                    channel)
                self._config2_svc = company_configs_pb2_grpc.CompanyConfigsServiceStub(
                    channel)
                for ct in self.config_types.values():
                    cls = [x for x in dir(ct.grpc) if x[-4:] == 'Stub']
                    if cls:
                        ct.svc = getattr(ct.grpc, cls[0])(channel)
            except:
                logging.exception('Error connecting to grpc service')
                time.sleep(30)

    async def connect(self):
        """Connect to the server and authenticate."""
        if self._cfg_lock is None:
            self._cfg_lock = asyncio.Lock()
        async with self._cfg_lock:
            await asyncio.get_event_loop().run_in_executor(self._executor,
                                                           self._connect)

    async def _run_subscription(self, sub, registry=None, key=None, slot=None):
        """Run a subscription to completion with guaranteed teardown.

        The single canonical lifecycle path for every subscription entry point:
        await sub.run(), then ALWAYS close it and DEREGISTER it -- on
        cancellation, a clean break, or a swallowed error alike -- so no dead
        subscription is ever left registered.  Deregistration covers both holders
        a subscription can live in: a `registry` dict under `key` (the cfg
        subscriptions) and/or a `slot` attribute name on this client (the nm_ent
        / nm_ref singletons).  Each is cleared only if it still points at THIS
        sub (a newer subscription may have replaced it).  CancelledError is
        suppressed: cancellation of an entry-point coroutine is a normal stop."""
        try:
            await sub.run()
        except asyncio.CancelledError:
            pass
        finally:
            sub.close()
            if registry is not None and registry.get(key) is sub:
                del registry[key]
            if slot is not None and getattr(self, slot, None) is sub:
                setattr(self, slot, None)

    async def nm_ent_subscribe(self, cbk, company_id=None, resolve=False, handler=None):
        """Block in this coroutine and call cbk as notifications arrive."""
        if not self._nm_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataSubscribeRequest(
            allow_ipv6_and_shorter_ipv4_prefixes=True)
        req.do_not_resolve_references = not resolve
        if company_id:
            req.company_id = company_id
        self._nm_ent_subscription = Subscription(self,
                                                 self._nm_svc.Subscribe,
                                                 req,
                                                 cbk,
                                                 handler=handler,
                                                 kind='nm_ent')
        await self._run_subscription(self._nm_ent_subscription,
                                     slot='_nm_ent_subscription')

    async def nm_ref_subscribe(self, cbk, company_id=None, handler=None):
        """Subscribe to Network Metadata Reference subscriptions."""
        if not self._nm_ref_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataReferenceSubscribeRequest()
        if company_id:
            req.company_id = company_id
        self._nm_ref_subscription = Subscription(self,
                                                 self._nm_ref_svc.Subscribe,
                                                 req,
                                                 cbk,
                                                 handler=handler,
                                                 kind='nm_ref')
        await self._run_subscription(self._nm_ref_subscription,
                                     slot='_nm_ref_subscription')

    async def nm_list(self, company_id=None):
        if not self._nm_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataListRequest()
        if company_id:
            req.company_id = company_id
        return await self._grpc_op(self._nm_svc.List, req)

    async def nm_set(self, entity, company_id=None):
        if not self._nm_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataSetRequest(entity=entity)
        if company_id:
            req.company_id = company_id
        return await self._grpc_op(self._nm_svc.Set, req)

    async def nm_del(self, entity, company_id=None):
        if not self._nm_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataDeleteRequest(entity=entity)
        if company_id:
            req.company_id = company_id
        return await self._grpc_op(self._nm_svc.Delete, req)

    async def nm_ref_list(self, company_id=None):
        if not self._nm_ref_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataReferenceListRequest()
        if company_id:
            req.company_id = company_id
        return await self._grpc_op(self._nm_ref_svc.List, req)

    async def nm_ref_set(self, ref, company_id=None):
        if not self._nm_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataReferenceSetRequest(entity=ref)
        if company_id:
            req.company_id = company_id
        return await self._grpc_op(self._nm_ref_svc.Set, req)

    async def nm_ref_del(self, ref, company_id=None):
        if not self._nm_svc:
            await self.connect()
        req = network_metadata_pb2.NetworkMetadataReferenceDeleteRequest(entity=ref)
        if company_id:
            req.company_id = company_id
        return await self._grpc_op(self._nm_ref_svc.Delete, req)

    async def push_flexes(self, req):
        if not self._metrics_svc:
            await self.connect()
        return await self._grpc_op(self._metrics_svc.PushFlexes, req, timeout=300)

    async def get_config(self, config_type, config_id, company_id=None):
        """Get a config instance by type and id"""
        ct = self.config_types.get(config_type)
        if ct is None:
            raise ValueError('Unsupported config type')
        if not ct.svc:
            await self.connect()
        req = ct.get_req()
        req.company_id = company_id
        setattr(req, f'{ct.prefix}_id', config_id)
        r = await self._grpc_op(ct.svc.Get, req)
        if bool(r.success):
            return getattr(r, ct.prefix)
        else:  # probably a race where it was deleted before we could get it
            logging.warning('Failed to get %s config "%s"', ct.name, config_id)
            return None

    async def set_config(self, config_type, obj, company_id=None):
        ct = self.config_types.get(config_type)
        if ct is None:
            raise ValueError('Unsupported config type')
        if not ct.svc:
            await self.connect()
        req = ct.set_req()
        req.company_id = company_id
        getattr(req, ct.prefix).CopyFrom(obj)
        r = await self._grpc_op(ct.svc.Set, req)
        return bool(r.success)

    async def del_config(self, config_type, config_id, company_id=None):
        ct = self.config_types.get(config_type)
        if ct is None:
            raise ValueError('Unsupported config type')
        if not ct.svc:
            await self.connect()
        req = ct.del_req()
        req.company_id = company_id
        setattr(req, f'{ct.prefix}_id', config_id)
        r = await self._grpc_op(ct.svc.Delete, req)
        return bool(r.success)

    async def sub_config(self, config_type, config_id, cbk, company_id=None):
        """
        Subscribe to notifications for updates to a config instance.
        The callback takes args (company_id, config_type, message).
        """

        async def _cbk(msg, _):
            nonlocal company_id, config_type, cbk
            await cbk(company_id, config_type, msg)

        ct = self.config_types.get(config_type)
        if not ct.svc:
            await self.connect()
        req = ct.sub_req()
        req.company_id = company_id
        setattr(req, f'{ct.prefix}_id', config_id)
        sub = self._cfg_subscriptions[config_id] = Subscription(self,
                                                                ct.svc.Subscribe,
                                                                req,
                                                                _cbk,
                                                                kind='cfg:{}'.format(
                                                                    ct.name))
        await self._run_subscription(sub,
                                     registry=self._cfg_subscriptions,
                                     key=config_id)

    async def sub_company_configs(self, cbk, company_id=None, subscribe_all=None):
        """Subscribe to company config updates."""
        if not self._config2_svc:
            await self.connect()
        req = company_configs_pb2.CompanyConfigSubscribeRequest()
        if company_id:
            req.company_id = company_id
        elif subscribe_all:
            req.subscribe_all = True
        sub = Subscription(self,
                           self._config2_svc.Subscribe,
                           req,
                           cbk,
                           kind='company_configs')
        await self._run_subscription(sub)
