decentral1se/chmxdx_letters.py

## chmxdx_letters.py
from contextlib import asynccontextmanager
from string import ascii_letters
from typing import Dict

import attr
from trio import (
    BrokenResourceError,
    ClosedResourceError,
    Lock,
    MemoryReceiveChannel,
    MemorySendChannel,
    open_memory_channel,
    open_nursery,
    run,
)
from trio.abc import Stream
from trio.testing import memory_stream_pair


@attr.s(auto_attribs=True)
class Channel:
    id: int

    protocol: "Protocol"
    stream: Stream

    sendq: MemorySendChannel
    recvq: MemoryReceiveChannel

    async def handler(self, msg: bytes):
        print(f"P{self.protocol.id}::C{self.id} got {msg.decode()}")

    async def asend(self, msg: bytes):
        await self.protocol.asend(msg)

    async def recv_task(self):
        while True:
            try:
                msg = await self.recvq.receive()
                await self.handler(msg)
                if msg == b"Z":
                    return
            except (ClosedResourceError, BrokenResourceError):
                print("P{self.protocol.id} memory channel down, exiting")
                return


@attr.s(auto_attribs=True)
class Message:
    id: int
    val: bytes


@attr.s(auto_attribs=True)
class Protocol:
    id: int
    lock: Lock
    stream: Stream
    channels: Dict[int, Channel] = attr.Factory(dict)
    BUFFER: int = 15
    MESSAGE_LENGTH: int = 2
    TIMEOUT: int = 10

    def parse(self, msg: bytes):
        id, val = bytes(msg[0:1]), bytes(msg[1:2])
        return Message(id=int(id), val=val)

    async def asend(self, msg: bytes):
        async with self.lock:
            await self.stream.send_all(msg)

    async def recv_task(self):
        while True:
            try:
                resp = await self.stream.receive_some(self.MESSAGE_LENGTH)
                msg = self.parse(resp)
                if msg.id not in self.channels:
                    continue
                channel = self.channels[msg.id]
                await channel.sendq.send(msg.val)
            except (ClosedResourceError, BrokenResourceError):
                print("P{self.id} stream down, exiting")
                return

    @asynccontextmanager
    async def channel(self, id: int):
        sendq, recvq = open_memory_channel(self.BUFFER)
        async with open_nursery() as nursery:
            channel = Channel(
                id=id,
                stream=self.stream,
                protocol=self,
                sendq=sendq,
                recvq=recvq,
            )
            nursery.start_soon(channel.recv_task)
            self.channels[id] = channel
            yield channel


async def channel(id, protocol):
    async with protocol.channel(id=id) as channel:
        for letter in ascii_letters:
            message = (b"%d" % id) + letter.encode()
            await channel.asend(message)


async def left_side_protocol(stream):
    async with open_nursery() as nursery:
        protocol = Protocol(id=1, stream=stream, lock=Lock())
        nursery.start_soon(protocol.recv_task)
        async with open_nursery() as n:
            for i in range(0, 5):
                n.start_soon(lambda: channel(id=i, protocol=protocol))
        nursery.cancel_scope.cancel()


async def right_side_protocol(stream):
    async with open_nursery() as nursery:
        protocol = Protocol(id=2, stream=stream, lock=Lock())
        nursery.start_soon(protocol.recv_task)
        async with open_nursery() as n:
            for i in range(0, 5):
                n.start_soon(lambda: channel(id=i, protocol=protocol))
        nursery.cancel_scope.cancel()


async def main():
    left_side_stream, right_side_stream = memory_stream_pair()

    async with open_nursery() as nursery:
        nursery.start_soon(left_side_protocol, left_side_stream)
        nursery.start_soon(right_side_protocol, right_side_stream)


run(main)
	from contextlib import asynccontextmanager
	from string import ascii_letters
	from typing import Dict

	import attr
	from trio import (
	BrokenResourceError,
	ClosedResourceError,
	Lock,
	MemoryReceiveChannel,
	MemorySendChannel,
	open_memory_channel,
	open_nursery,
	run,
	)
	from trio.abc import Stream
	from trio.testing import memory_stream_pair


	@attr.s(auto_attribs=True)
	class Channel:
	id: int

	protocol: "Protocol"
	stream: Stream

	sendq: MemorySendChannel
	recvq: MemoryReceiveChannel

	async def handler(self, msg: bytes):
	print(f"P{self.protocol.id}::C{self.id} got {msg.decode()}")

	async def asend(self, msg: bytes):
	await self.protocol.asend(msg)

	async def recv_task(self):
	while True:
	try:
	msg = await self.recvq.receive()
	await self.handler(msg)
	if msg == b"Z":
	return
	except (ClosedResourceError, BrokenResourceError):
	print("P{self.protocol.id} memory channel down, exiting")
	return


	@attr.s(auto_attribs=True)
	class Message:
	id: int
	val: bytes


	@attr.s(auto_attribs=True)
	class Protocol:
	id: int
	lock: Lock
	stream: Stream
	channels: Dict[int, Channel] = attr.Factory(dict)
	BUFFER: int = 15
	MESSAGE_LENGTH: int = 2
	TIMEOUT: int = 10

	def parse(self, msg: bytes):
	id, val = bytes(msg[0:1]), bytes(msg[1:2])
	return Message(id=int(id), val=val)

	async def asend(self, msg: bytes):
	async with self.lock:
	await self.stream.send_all(msg)

	async def recv_task(self):
	while True:
	try:
	resp = await self.stream.receive_some(self.MESSAGE_LENGTH)
	msg = self.parse(resp)
	if msg.id not in self.channels:
	continue
	channel = self.channels[msg.id]
	await channel.sendq.send(msg.val)
	except (ClosedResourceError, BrokenResourceError):
	print("P{self.id} stream down, exiting")
	return

	@asynccontextmanager
	async def channel(self, id: int):
	sendq, recvq = open_memory_channel(self.BUFFER)
	async with open_nursery() as nursery:
	channel = Channel(
	id=id,
	stream=self.stream,
	protocol=self,
	sendq=sendq,
	recvq=recvq,
	)
	nursery.start_soon(channel.recv_task)
	self.channels[id] = channel
	yield channel


	async def channel(id, protocol):
	async with protocol.channel(id=id) as channel:
	for letter in ascii_letters:
	message = (b"%d" % id) + letter.encode()
	await channel.asend(message)


	async def left_side_protocol(stream):
	async with open_nursery() as nursery:
	protocol = Protocol(id=1, stream=stream, lock=Lock())
	nursery.start_soon(protocol.recv_task)
	async with open_nursery() as n:
	for i in range(0, 5):
	n.start_soon(lambda: channel(id=i, protocol=protocol))
	nursery.cancel_scope.cancel()


	async def right_side_protocol(stream):
	async with open_nursery() as nursery:
	protocol = Protocol(id=2, stream=stream, lock=Lock())
	nursery.start_soon(protocol.recv_task)
	async with open_nursery() as n:
	for i in range(0, 5):
	n.start_soon(lambda: channel(id=i, protocol=protocol))
	nursery.cancel_scope.cancel()


	async def main():
	left_side_stream, right_side_stream = memory_stream_pair()

	async with open_nursery() as nursery:
	nursery.start_soon(left_side_protocol, left_side_stream)
	nursery.start_soon(right_side_protocol, right_side_stream)


	run(main)