solana-py

How to create a websocket manager in Solana-py


I am using Python v3.7 + solana v0.21.0 and trying to create a websocket manager handling several different subscriptions, but the code seems to block whenever websocket.recv() or asyncstdlib.enumerate(websocket) is called. Examples from the docs:

import asyncio
from asyncstdlib import enumerate
from solana.rpc.websocket_api import connect


async def main():
    # First example using websocket.recv()
    async with connect("wss://api.devnet.solana.com") as websocket:
        await websocket.logs_subscribe()
        first_resp = await websocket.recv()
        subscription_id = first_resp.result
        next_resp = await websocket.recv()         <--- Blocks until msg received.
        print(next_resp)
        await websocket.logs_unsubscribe(subscription_id)


    # Second example using client as an infinite asynchronous iterator:
    async with connect("wss://api.devnet.solana.com") as websocket:  
        await websocket.logs_subscribe()
        first_resp = await websocket.recv()
        subscription_id = first_resp.result
        async for idx, msg in enumerate(websocket): <--- Blocks until msg received.
            if idx == 3:
                break
            print(msg)
        await websocket.logs_unsubscribe(subscription_id)

asyncio.run(main())

The idea is being able to iterate an infinite loop so new subscriptions can be added to the websocket for instance:

from solana.rpc.request_builder import LogsSubscribeFilter
from solana.rpc.websocket_api import connect
from asgiref.sync import sync_to_async
from solana.publickey import PublicKey
from time import sleep


async def websocket_manager(rpc: str):
  async with connect(rpc) as websocket:
        while True:
            active_pubkeys = await sync_to_async(get_my_active_pubkeys)()
            if active_pubkeys:
                # Add missing pubkeys
                for pubkey in active_pubkeys:
                    if ws.get(pubkey) in websocket.subscriptions:
                        continue

                    print(f"Subscribe to {pubkey}")
                    await websocket.logs_subscribe(LogsSubscribeFilter.mentions(PublicKey(pubkey)))
                    first_resp = await websocket.recv()
                    ws[pubkey] = first_resp.result # Maps the pubkey to the subscription ID

                # Delete non used subscriptions:
                for non_used_pubkey in set(active_pubkeys) ^ set(ws.keys()):
                    if non_used_pubkey in ws:
                        print(f"Delete subscription for pubkey #{non_used_pubkey}")
                        websocket.account_unsubscribe(ws[non_used_pubkey])
                        ws.pop(non_used_pubkey)

            # <-- HERE HOW TO ITERATE SUBSCRIPTIONS WITHOUT BLOCKING THE MANAGER????

            sleep(30)  # Sleep for 30 seconds

Would be it be safe using a new thread or a subprocess for reading the websocket messages so they don't block the main function?


Solution

  • I figured it out, just needed to use task = asyncio.create_task(my_msg_listener_function(websocket, ws)) to have another task running concurrently and taking care of the websocket messages.