53def run(coro: Coroutine[Any, Any, None]) -> None:
54    """`finsy.run` provides a useful wrapper around `asyncio.run`.
55
56    This function implements common boilerplate for running a Finsy application.
57
58    - Set up basic logging to stderr at the INFO log level.
59    - Set up a signal handler for SIGTERM that shuts down gracefully.
60    - Set up caching of P4Info data so common definitions are re-used.
61
62    Example:
63
64    ```python
65    import finsy as fy
66
67    async def main():
68        async with fy.Switch("sw", "127.0.0.1:50001") as sw:
69            print(sw.p4info)
70
71    if __name__ == "__main__":
72        fy.run(main())
73    ```
74
75    If you choose to use `asyncio.run` instead, your P4Schema/P4Info objects
76    will not be eligible for sharing. You can create your own `P4SchemaCache`
77    context manager to implement this.
78    """
79    try:
80        asyncio.run(_finsy_boilerplate(coro))
81    except (KeyboardInterrupt, asyncio.CancelledError):
82        pass

 30@final
 31class Controller:
 32    """Represents a collection of P4Runtime switches.
 33
 34    Each `Switch` in the Controller is identified by its name. Each name must
 35    be unique.
 36
 37    ```
 38    switches = [
 39        fy.Switch("sw1", "10.0.0.1:50000"),
 40        fy.Switch("sw2", "10.0.0.2:50000"),
 41    ]
 42    controller = fy.Controller(switches)
 43    await controller.run()
 44    ```
 45
 46    A Controller can be running or stopped. There are two ways to run a
 47    Controller. You can use the `Controller.run()` method, or you can use
 48    a Controller as a context manager.
 49
 50    ```
 51    controller = fy.Controller(switches)
 52    async with controller:
 53        # Let controller run for 60 seconds.
 54        await asyncio.sleep(60)
 55    ```
 56
 57    You can `add` or `remove` switches regardless of whether a Controller is
 58    running or not. If the Controller is not running, adding or removing a Switch is
 59    instantaneous. If the Controller is running, adding a Switch will start
 60    it running asynchronously. Removing a Switch will schedule the Switch to stop,
 61    but defer actual removal until the Switch has stopped asynchronously.
 62
 63    When a switch starts inside a Controller, it fires the `CONTROLLER_ENTER`
 64    event. When a switch stops inside a Controller, it fires the `CONTROLLER_LEAVE`
 65    event.
 66
 67    A Controller supports these methods to access its contents:
 68
 69    - len(controller): Return number of switches in the controller.
 70    - controller[name]: Return the switch with the given name.
 71    - controller.get(name): Return the switch with the given name, or None if not found.
 72
 73    You can iterate over the switches in a Controller using a for loop:
 74
 75    ```
 76    for switch in controller:
 77        print(switch.name)
 78    ```
 79
 80    Any task or sub-task running inside a controller can retrieve its
 81    Controller object using the `Controller.current()` method.
 82    """
 83
 84    _switches: dict[str, Switch]
 85    _pending_removal: set[Switch]
 86    _task_count: CountdownFuture
 87
 88    _control_task: asyncio.Task[Any] | None = None
 89    "Keep track of controller's main task."
 90
 91    def __init__(self, switches: Iterable[Switch] = ()):
 92        """Initialize Controller object with an initial list of switches.
 93
 94        Args:
 95            switches: a collection or other iterable of Switch objects.
 96        """
 97        self._switches = {}
 98        self._pending_removal = set()
 99        self._task_count = CountdownFuture()
100
101        for switch in switches:
102            if switch.name in self._switches:
103                raise ValueError(f"Switch named {switch.name!r} already exists")
104            self._switches[switch.name] = switch
105
106    @property
107    def running(self) -> bool:
108        "True if Controller is running."
109        return self._control_task is not None
110
111    async def run(self) -> None:
112        "Run the controller."
113        async with self:
114            await wait_for_cancel()
115
116    def stop(self) -> None:
117        "Stop the controller if it is running."
118        if self._control_task is not None:
119            self._control_task.cancel()
120
121    async def __aenter__(self) -> Self:
122        "Run the controller as a context manager (see also run())."
123        assert not self.running, "Controller.__aenter__ is not re-entrant"
124        assert self._task_count.value() == 0
125        assert not self._pending_removal
126
127        self._control_task = asyncio.current_task()
128        _CONTROLLER.set(self)
129
130        try:
131            # Start each switch running.
132            for switch in self:
133                self._start_switch(switch)
134        except Exception:
135            self._control_task = None
136            _CONTROLLER.set(None)
137            raise
138
139        return self
140
141    async def __aexit__(
142        self,
143        _exc_type: type[BaseException] | None,
144        _exc_val: BaseException | None,
145        _exc_tb: TracebackType | None,
146    ) -> bool | None:
147        "Run the controller as a context manager (see also run())."
148        assert self.running
149
150        try:
151            # Stop all the switches.
152            for switch in self:
153                self._stop_switch(switch)
154
155            # Wait for switch tasks to finish.
156            await self._task_count.wait()
157
158        finally:
159            self._control_task = None
160            _CONTROLLER.set(None)
161
162    def add(self, switch: Switch) -> None:
163        """Add a switch to the controller.
164
165        If the controller is running, tell the switch to start.
166
167        Args:
168            switch: the Switch object.
169        """
170        if switch.name in self._switches:
171            raise ValueError(f"Switch named {switch.name!r} already exists")
172
173        self._switches[switch.name] = switch
174        if self.running:
175            self._start_switch(switch)
176
177    def remove(self, switch: Switch) -> asyncio.Event:
178        """Remove a switch from the controller.
179
180        If the controller is running, tell the switch to stop and schedule it
181        for removal when it fully stops.
182
183        Args:
184            switch: the Switch object.
185        """
186        name = switch.name
187        if self._switches.get(name, None) is not switch:
188            raise ValueError(f"Switch named {name!r} not found")
189
190        del self._switches[name]
191
192        event = asyncio.Event()
193        if self.running:
194            # When controller is running, event will complete when switch
195            # is actually stopped.
196            self._stop_switch(switch)
197            self._pending_removal.add(switch)
198
199            def _controller_leave(sw: Switch):
200                self._pending_removal.discard(sw)
201                event.set()
202
203            switch.ee.once(SwitchEvent.CONTROLLER_LEAVE, _controller_leave)  # type: ignore
204        else:
205            # When controller is not running, event completes immediately.
206            event.set()
207
208        return event
209
210    def _start_switch(self, switch: Switch):
211        "Start the switch's control task."
212        LOGGER.debug("Controller._start_switch: %r", switch)
213        assert switch._control_task is None  # pyright: ignore[reportPrivateUsage]
214
215        switch.ee.emit(SwitchEvent.CONTROLLER_ENTER, switch)
216
217        task = asyncio.create_task(switch.run(), name=f"fy:{switch.name}")
218        switch._control_task = task  # pyright: ignore[reportPrivateUsage]
219        self._task_count.increment()
220
221        def _switch_done(done: asyncio.Task[Any]):
222            switch._control_task = None  # pyright: ignore[reportPrivateUsage]
223            switch.ee.emit(SwitchEvent.CONTROLLER_LEAVE, switch)
224            self._task_count.decrement()
225
226            if not done.cancelled():
227                ex = done.exception()
228                if ex is not None:
229                    if not isinstance(ex, SwitchFailFastError):
230                        # The `fail_fast` error has already been logged. If
231                        # it's any other error, log it. (There shouldn't be
232                        # any other error.)
233                        LOGGER.critical(
234                            "Controller task %r failed",
235                            done.get_name(),
236                            exc_info=ex,
237                        )
238                    # Shutdown the program cleanly due to switch failure.
239                    raise SystemExit(99)
240
241        task.add_done_callback(_switch_done)
242
243    def _stop_switch(self, switch: Switch):
244        "Stop the switch's control task."
245        LOGGER.debug("Controller._stop_switch: %r", switch)
246
247        if switch._control_task is not None:  # pyright: ignore[reportPrivateUsage]
248            switch._control_task.cancel()  # pyright: ignore[reportPrivateUsage]
249
250    def __len__(self) -> int:
251        "Return the number of switches."
252        return len(self._switches)
253
254    def __iter__(self) -> Iterator[Switch]:
255        "Iterate over the switches."
256        return iter(self._switches.values())
257
258    def __getitem__(self, name: str) -> Switch:
259        "Retrieve a switch by name."
260        return self._switches[name]
261
262    def get(self, name: str) -> Switch | None:
263        "Retrieve a switch by name, or return None if not found."
264        return self._switches.get(name)
265
266    @staticmethod
267    def current() -> "Controller":
268        "Return the current Controller object."
269        result = _CONTROLLER.get()
270        if result is None:
271            raise RuntimeError("controller does not exist")
272        return result

68class LoggerAdapter(_BaseLoggerAdapter):
69    """Custom log adapter to include the name of the current task."""
70
71    def process(
72        self,
73        msg: Any,
74        kwargs: MutableMapping[str, Any],
75    ) -> tuple[Any, MutableMapping[str, Any]]:
76        """Process the logging message and keyword arguments passed in to a
77        logging call to insert contextual information.
78        """
79        task_name = _get_current_task_name()
80        return f"[{task_name}] {msg}", kwargs
81
82    def info(self, msg: Any, *args: Any, **kwargs: Any) -> None:
83        """INFO level uses a concise task name representation for readability."""
84        if self.logger.isEnabledFor(logging.INFO):
85            task_name = _get_current_task_name(True)
86            self.logger.info(f"[{task_name}] {msg}", *args, **kwargs)

 24@functools.total_ordering
 25class MACAddress:
 26    """Concrete class for a MAC address."""
 27
 28    __slots__ = ("_mac", "__weakref__")
 29    _mac: int
 30
 31    def __init__(self, value: object) -> None:
 32        "Construct a MAC address from a string, integer or bytes object."
 33        match value:
 34            case int():
 35                self._mac = _from_int(value)
 36            case bytes():
 37                self._mac = _from_bytes(value)
 38            case MACAddress():
 39                self._mac = value._mac
 40            case _:
 41                # Input argument is a MAC string, or an object that is
 42                # formatted as MAC string (behave like ipaddress module).
 43                self._mac = _from_string(str(value))
 44
 45    @property
 46    def packed(self) -> bytes:
 47        "Return the MAC address a byte string."
 48        return _to_bytes(self._mac)
 49
 50    @property
 51    def max_prefixlen(self) -> int:
 52        "Return the maximum prefix length (48 bits)."
 53        return _BIT_WIDTH
 54
 55    @property
 56    def is_multicast(self) -> bool:
 57        "Return true if MAC address has the multicast bit set."
 58        return self._mac & (1 << 40) != 0
 59
 60    @property
 61    def is_private(self) -> bool:
 62        "Return true if MAC address has the locally administered bit set."
 63        return self._mac & (1 << 41) != 0
 64
 65    @property
 66    def is_global(self) -> bool:
 67        "Return true if the locally administered bit is not set."
 68        return not self.is_private
 69
 70    @property
 71    def is_unspecified(self) -> bool:
 72        "Return true if MAC address is all zeros."
 73        return self._mac == 0
 74
 75    @property
 76    def is_broadcast(self) -> bool:
 77        "Return true if MAC address is the broadcast address."
 78        return self._mac == (1 << _BIT_WIDTH) - 1
 79
 80    def __int__(self) -> int:
 81        return self._mac
 82
 83    def __eq__(self, rhs: object) -> bool:
 84        if not isinstance(rhs, MACAddress):
 85            return NotImplemented
 86        return self._mac == rhs._mac
 87
 88    def __lt__(self, rhs: object) -> bool:
 89        if not isinstance(rhs, MACAddress):
 90            return NotImplemented
 91        return self._mac < rhs._mac
 92
 93    def __repr__(self) -> str:
 94        return f"MACAddress({_to_string(self._mac)!r})"
 95
 96    def __str__(self) -> str:
 97        return _to_string(self._mac)
 98
 99    def __hash__(self) -> int:
100        return hash(hex(self._mac))

1458@decodable("action_profile_group")
1459@dataclass(slots=True)
1460class P4ActionProfileGroup(_P4Writable):
1461    "Represents a P4Runtime ActionProfileGroup."
1462
1463    action_profile_id: str = ""
1464    _: KW_ONLY
1465    group_id: int = 0
1466    max_size: int = 0
1467    members: Sequence[P4Member] | None = None
1468
1469    def encode(self, schema: P4Schema) -> p4r.Entity:
1470        "Encode P4ActionProfileGroup as protobuf."
1471        if not self.action_profile_id:
1472            return p4r.Entity(action_profile_group=p4r.ActionProfileGroup())
1473
1474        profile = schema.action_profiles[self.action_profile_id]
1475
1476        if self.members is not None:
1477            members = [member.encode() for member in self.members]
1478        else:
1479            members = None
1480
1481        entry = p4r.ActionProfileGroup(
1482            action_profile_id=profile.id,
1483            group_id=self.group_id,
1484            members=members,
1485            max_size=self.max_size,
1486        )
1487        return p4r.Entity(action_profile_group=entry)
1488
1489    @classmethod
1490    def decode(cls, msg: p4r.Entity, schema: P4Schema) -> Self:
1491        "Decode protobuf to ActionProfileGroup data."
1492        entry = msg.action_profile_group
1493        if entry.action_profile_id == 0:
1494            return cls()
1495
1496        profile = schema.action_profiles[entry.action_profile_id]
1497
1498        if entry.members:
1499            members = [P4Member.decode(member) for member in entry.members]
1500        else:
1501            members = None
1502
1503        return cls(
1504            action_profile_id=profile.alias,
1505            group_id=entry.group_id,
1506            max_size=entry.max_size,
1507            members=members,
1508        )
1509
1510    def action_str(self, _schema: P4Schema) -> str:
1511        "Return string representation of the weighted members."
1512        if not self.members:
1513            return ""
1514
1515        return " ".join(
1516            [f"{member.weight}*{member.member_id:#x}" for member in self.members]
1517        )

1356@decodable("action_profile_member")
1357@dataclass(slots=True)
1358class P4ActionProfileMember(_P4Writable):
1359    "Represents a P4Runtime ActionProfileMember."
1360
1361    action_profile_id: str = ""
1362    _: KW_ONLY
1363    member_id: int = 0
1364    action: P4TableAction | None = None
1365
1366    def encode(self, schema: P4Schema) -> p4r.Entity:
1367        "Encode P4ActionProfileMember as protobuf."
1368        if not self.action_profile_id:
1369            return p4r.Entity(action_profile_member=p4r.ActionProfileMember())
1370
1371        profile = schema.action_profiles[self.action_profile_id]
1372
1373        if self.action:
1374            action = self.action.encode_action(schema)
1375        else:
1376            action = None
1377
1378        entry = p4r.ActionProfileMember(
1379            action_profile_id=profile.id,
1380            member_id=self.member_id,
1381            action=action,
1382        )
1383        return p4r.Entity(action_profile_member=entry)
1384
1385    @classmethod
1386    def decode(cls, msg: p4r.Entity, schema: P4Schema) -> Self:
1387        "Decode protobuf to ActionProfileMember data."
1388        entry = msg.action_profile_member
1389        if entry.action_profile_id == 0:
1390            return cls()
1391
1392        profile = schema.action_profiles[entry.action_profile_id]
1393
1394        if entry.HasField("action"):
1395            action = P4TableAction.decode_action(entry.action, schema)
1396        else:
1397            action = None
1398
1399        return cls(
1400            action_profile_id=profile.alias,
1401            member_id=entry.member_id,
1402            action=action,
1403        )
1404
1405    def action_str(self, schema: P4Schema) -> str:
1406        "Format the action as a human-readable, canonical string."
1407        if self.action is None:
1408            return NOACTION_STR
1409        return self.action.format_str(schema)

239class P4Client:
240    "Implements a P4Runtime client."
241
242    _address: str
243    _credentials: GRPCCredentialsTLS | None
244    _wait_for_ready: bool
245    _channel: grpc.aio.Channel | None = None
246    _stub: p4r_grpc.P4RuntimeStub | None = None
247    _stream: _P4StreamTypeAlias | None = None
248    _complete_request: Callable[[pbutil.PBMessage], None] | None = None
249
250    _schema: P4Schema | None = None
251    "Annotate log messages using this optional P4Info schema."
252
253    def __init__(
254        self,
255        address: str,
256        credentials: GRPCCredentialsTLS | None = None,
257        *,
258        wait_for_ready: bool = True,
259    ) -> None:
260        self._address = address
261        self._credentials = credentials
262        self._wait_for_ready = wait_for_ready
263
264    @property
265    def channel(self) -> grpc.aio.Channel | None:
266        "Return the GRPC channel object, or None if the channel is not open."
267        return self._channel
268
269    async def __aenter__(self) -> Self:
270        await self.open()
271        return self
272
273    async def __aexit__(self, *args: Any) -> bool | None:
274        await self.close()
275
276    async def open(
277        self,
278        *,
279        schema: P4Schema | None = None,
280        complete_request: Callable[[pbutil.PBMessage], None] | None = None,
281    ) -> None:
282        """Open the client channel.
283
284        Note: This method is `async` for forward-compatible reasons.
285        """
286        assert self._channel is None
287        assert self._stub is None
288
289        # Increase max_metadata_size from 8 KB to 64 KB.
290        options = GRPCOptions(
291            max_metadata_size=64 * 1024,  # 64 kilobytes
292            max_reconnect_backoff_ms=15000,  # 15.0 seconds
293        )
294
295        self._channel = grpc_channel(
296            self._address,
297            credentials=self._credentials,
298            options=options,
299            client_type="P4Client",
300        )
301
302        self._stub = p4r_grpc.P4RuntimeStub(self._channel)
303        self._schema = schema
304        self._complete_request = complete_request
305
306    async def close(self) -> None:
307        "Close the client channel."
308        if self._channel is not None:
309            LOGGER.debug("P4Client: close channel %r", self._address)
310
311            if self._stream is not None:
312                self._stream.cancel()
313                self._stream = None
314
315            await self._channel.close()
316            self._channel = None
317            self._stub = None
318            self._schema = None
319            self._complete_request = None
320
321    async def send(self, msg: p4r.StreamMessageRequest) -> None:
322        """Send a message to the stream."""
323        assert self._stub is not None
324
325        if not self._stream or self._stream.done():
326            self._stream = cast(
327                _P4StreamTypeAlias,
328                self._stub.StreamChannel(wait_for_ready=self._wait_for_ready),  # type: ignore
329            )
330
331        self._log_msg(msg)
332
333        try:
334            await self._stream.write(msg)
335        except grpc.RpcError as ex:
336            raise P4ClientError(ex, "send") from None
337
338    async def receive(self) -> p4r.StreamMessageResponse:
339        """Read a message from the stream."""
340        assert self._stream is not None
341
342        try:
343            msg = cast(
344                p4r.StreamMessageResponse,
345                await self._stream.read(),  # type: ignore
346            )
347            if msg is GRPC_EOF:
348                # Treat EOF as a protocol violation.
349                raise RuntimeError("P4Client.receive got EOF!")
350
351        except grpc.RpcError as ex:
352            raise P4ClientError(ex, "receive") from None
353
354        self._log_msg(msg)
355        return msg
356
357    @overload
358    async def request(
359        self, msg: p4r.WriteRequest
360    ) -> p4r.WriteResponse: ...  # pragma: no cover
361
362    @overload
363    async def request(
364        self, msg: p4r.GetForwardingPipelineConfigRequest
365    ) -> p4r.GetForwardingPipelineConfigResponse: ...  # pragma: no cover
366
367    @overload
368    async def request(
369        self, msg: p4r.SetForwardingPipelineConfigRequest
370    ) -> p4r.SetForwardingPipelineConfigResponse: ...  # pragma: no cover
371
372    @overload
373    async def request(
374        self, msg: p4r.CapabilitiesRequest
375    ) -> p4r.CapabilitiesResponse: ...  # pragma: no cover
376
377    async def request(self, msg: pbutil.PBMessage) -> pbutil.PBMessage:
378        "Send a unary-unary P4Runtime request and wait for the response."
379        assert self._stub is not None
380
381        if self._complete_request:
382            self._complete_request(msg)
383
384        msg_type = type(msg).__name__
385        assert msg_type.endswith("Request")
386        rpc_method = getattr(self._stub, msg_type[:-7])
387
388        self._log_msg(msg)
389        try:
390            reply = await rpc_method(
391                msg,
392                timeout=_DEFAULT_RPC_TIMEOUT,
393            )
394        except grpc.RpcError as ex:
395            raise P4ClientError(ex, msg_type, msg=msg, schema=self._schema) from None
396
397        self._log_msg(reply)
398        return reply
399
400    async def request_iter(
401        self, msg: p4r.ReadRequest
402    ) -> AsyncIterator[p4r.ReadResponse]:
403        "Send a unary-stream P4Runtime read request and wait for the responses."
404        assert self._stub is not None
405
406        if self._complete_request:
407            self._complete_request(msg)
408
409        msg_type = type(msg).__name__
410        assert msg_type.endswith("Request")
411        rpc_method = getattr(self._stub, msg_type[:-7])
412
413        self._log_msg(msg)
414        try:
415            async for reply in rpc_method(
416                msg,
417                timeout=_DEFAULT_RPC_TIMEOUT,
418            ):
419                self._log_msg(reply)
420                yield reply
421        except grpc.RpcError as ex:
422            raise P4ClientError(ex, msg_type) from None
423
424    def _log_msg(self, msg: pbutil.PBMessage) -> None:
425        "Log a P4Runtime request or response."
426        pbutil.log_msg(self._channel, msg, self._schema)