|
4 | 4 | import argparse
|
5 | 5 | import asyncio
|
6 | 6 | import contextlib
|
| 7 | +from contextvars import ContextVar |
7 | 8 | import enum
|
8 | 9 | import os
|
9 | 10 | import pathlib
|
@@ -1529,8 +1530,45 @@ def print_version(self):
|
1529 | 1530 | print(labgrid_version())
|
1530 | 1531 |
|
1531 | 1532 |
|
1532 |
| -def start_session(address, extra, debug=False): |
1533 |
| - loop = asyncio.get_event_loop() |
| 1533 | +_loop: ContextVar["asyncio.AbstractEventLoop | None"] = ContextVar("_loop", default=None) |
| 1534 | + |
| 1535 | + |
| 1536 | +def ensure_event_loop(external_loop=None): |
| 1537 | + """Get the event loop for this thread, or create a new event loop.""" |
| 1538 | + # get stashed loop |
| 1539 | + loop = _loop.get() |
| 1540 | + |
| 1541 | + # ignore closed stashed loop |
| 1542 | + if loop and loop.is_closed(): |
| 1543 | + loop = None |
| 1544 | + |
| 1545 | + if external_loop: |
| 1546 | + # if a loop is stashed, expect it to be the same as the external one |
| 1547 | + if loop: |
| 1548 | + assert loop is external_loop |
| 1549 | + _loop.set(external_loop) |
| 1550 | + return external_loop |
| 1551 | + |
| 1552 | + # return stashed loop |
| 1553 | + if loop: |
| 1554 | + return loop |
| 1555 | + |
| 1556 | + try: |
| 1557 | + # if called from async code, try to get current's thread loop |
| 1558 | + loop = asyncio.get_running_loop() |
| 1559 | + except RuntimeError: |
| 1560 | + # no previous, external or running loop found, create a new one |
| 1561 | + loop = asyncio.new_event_loop() |
| 1562 | + asyncio.set_event_loop(loop) |
| 1563 | + |
| 1564 | + # stash it |
| 1565 | + _loop.set(loop) |
| 1566 | + return loop |
| 1567 | + |
| 1568 | + |
| 1569 | +def start_session(address, extra, debug=False, loop=None): |
| 1570 | + loop = ensure_event_loop(loop) |
| 1571 | + |
1534 | 1572 | if debug:
|
1535 | 1573 | loop.set_debug(True)
|
1536 | 1574 |
|
@@ -2040,7 +2078,9 @@ def main():
|
2040 | 2078 | coordinator_address = os.environ.get("LG_COORDINATOR", "127.0.0.1:20408")
|
2041 | 2079 |
|
2042 | 2080 | logging.debug('Starting session with "%s"', coordinator_address)
|
2043 |
| - session = start_session(coordinator_address, extra, args.debug) |
| 2081 | + loop = asyncio.new_event_loop() |
| 2082 | + asyncio.set_event_loop(loop) |
| 2083 | + session = start_session(coordinator_address, extra=extra, debug=args.debug, loop=loop) |
2044 | 2084 | logging.debug("Started session")
|
2045 | 2085 |
|
2046 | 2086 | try:
|
|
0 commit comments