One socket, many requests

# -*- coding: utf-8 -*-

# stdlib
import itertools
import socket
import threading

# Zato
from zato.common.sdk import Connector, ConnectionLost, Field

class PaymentsClient:
    """ A client for a payment switch that multiplexes many in-flight requests over one persistent
    socket - each request has a correlation ID, replies can arrive in any order and a reader
    loop matches them back to the calls that wait for them (ISO 8583, FIX style).
    """
    def __init__(self, host:'str', port:'int') -> 'None':

        # The one socket all the requests share.
        self.socket = socket.create_connection((host, port))
        self.reader_file = self.socket.makefile('r', encoding='utf8')

        # Guards the pending map and writes to the shared socket.
        self.lock = threading.Lock()

        # Calls in flight, keyed by their correlation IDs.
        self.pending = {}

        # Where the correlation IDs come from.
        self.counter = itertools.count(1)

        # Set to False by the reader loop once the socket is gone.
        self.is_connected = True

        # The reader loop matches replies back to waiting calls for as long as the connection lives.
        reader_thread = threading.Thread(target=self._read_loop, daemon=True)
        reader_thread.start()

    def _read_loop(self) -> 'None':

        for line in self.reader_file:
            text = line.strip()
            corr_id, _, payload = text.partition(' ')

            # The call this reply belongs to may have given up already, e.g. it timed out.
            with self.lock:
                holder = self.pending.pop(corr_id, None)

            if holder:
                holder['response'] = payload
                holder['event'].set()

        # The loop ended, which means the socket is gone - wake up everyone still waiting.
        self.is_connected = False

        with self.lock:
            for holder in self.pending.values():
                holder['event'].set()
            self.pending.clear()

    def request(self, payload:'str') -> 'str':

        # A dead socket cannot deliver anything - the framework will reconnect.
        if not self.is_connected:
            raise ConnectionLost('The payment switch connection is down')

        corr_id = str(next(self.counter))
        holder = {'event': threading.Event(), 'response': None}

        # Register the call and send its request under one lock, so the reply cannot
        # arrive before the call is registered.
        with self.lock:
            self.pending[corr_id] = holder
            self.socket.sendall(f'{corr_id} {payload}\n'.encode('utf8'))

        # Wait for the reader loop to match the reply back to this call.
        _ = holder['event'].wait()

        # A wake-up without a response means the connection died while this call was in flight.
        if holder['response'] is None:
            raise ConnectionLost('The payment switch connection went down mid-call')

        return holder['response']

    def close(self) -> 'None':
        self.reader_file.close()
        self.socket.close()

class PaymentsConnector(Connector):
    """ Wraps the payment switch as a connection type that services access through self.out.payments.
    """
    type = 'payments'

    # Configuration schema
    host = Field.Text()
    port = Field.Int(default=9970)

    def create_client(self) -> 'PaymentsClient':
        return PaymentsClient(self.config.host, self.config.port)

    def ping(self, client:'PaymentsClient') -> 'None':
        client.request('ping')

    def on_stop(self, client:'PaymentsClient') -> 'None':
        client.close()

    def authorize(self, payload:'str') -> 'str':
        return self.client.request(payload)

# -*- coding: utf-8 -*-

# Zato
from zato.server.service import Service

class AuthorizePayment(Service):
    """ Authorizes one payment over the multiplexed connection to the payment switch.
    """
    name = 'demo.payments.authorize'

    input = 'payload'

    def handle(self) -> 'None':
        conn = self.out.payments['My Payments']
        response = conn.authorize(self.request.input.payload)
        self.response.payload = response