Feature: Allow concurrent RPC calls on rabbit implementation about faststream HOT 5 OPEN

from asyncio import Future
from contextlib import asynccontextmanager
from typing import Annotated, Any
from uuid import uuid4

from faststream import Context, FastStream, context
from faststream.rabbit import RabbitBroker


@asynccontextmanager
async def lifespan():
    context.set_global("replies_container", {})
    yield


RepliesContainer = Annotated[dict[str, Future[Any]], Context("replies_container")]

broker = RabbitBroker()
app = FastStream(broker, lifespan=lifespan)


@broker.subscriber("replies")
def handle_reply(
    data,
    container: RepliesContainer,
    cor_id: str = Context("message.correlation_id"),
):
    if (future := container.pop(cor_id, None)) and not future.done():
        future.set_result(data)


async def custom_rpc(
    broker: NatsBroker,
    msg: Any,
    subject: str,
    container: RepliesContainer,
) -> Any:
    cor_id = uuid4().hex
    container[cor_id] = result_future = Future()
    await broker.publish(msg, subject, correlation_id=cor_id)
    return await result_future


# Emulation

@broker.subscriber("in")
@broker.publisher("replies")
async def handle_request(data):
    """Your hardcoded not-FastStream service."""
    return data


@app.after_startup
async def t(container: RepliesContainer):
    response = await custom_rpc(broker, "test", "in", container)
    print(response)

import os
import random

from faststream.rabbit import RabbitBroker
from faststream import Context, FastStream, context
from fastapi import FastAPI
import uvicorn

from contextlib import asynccontextmanager
from asyncio import Future
from typing import Annotated, Any
import uuid


RepliesContainer = Annotated[dict[str, Future[Any]], Context("replies_container")]
broker = RabbitBroker()


@asynccontextmanager
async def fastapi_lifespan(fastapi_app: FastAPI):
    context.set_global("replies_container", {})
    await broker.start()
    yield
    await broker.close()


fastapi_app = FastAPI(lifespan=fastapi_lifespan)


@broker.subscriber("replies")
def handle_reply(
    data,
    container: RepliesContainer,
    cor_id: str = Context("message.correlation_id"),
):
    print('handle_reply receives message', data, cor_id, os.getpid())
    if (future := container.pop(cor_id, None)) and not future.done():
        future.set_result(data)


async def custom_rpc(
    broker: RabbitBroker,
    msg: Any,
    subject: str,
    container: RepliesContainer,
) -> Any:
    cor_id = uuid.uuid4().hex
    container[cor_id] = result_future = Future()
    print('Publish task', msg, cor_id, os.getpid())
    await broker.publish(msg, subject, correlation_id=cor_id)
    print('Publish task success', msg, cor_id, os.getpid())
    return await result_future


@fastapi_app.post('/push_task')
async def push_task():
    msg = 'test' + str(random.randint(100, 999))
    container = context.get('replies_container')
    response = await custom_rpc(broker, msg, "in", container)
    print('response received', msg, response, os.getpid())

from faststream import Context, FastStream, context
from faststream.rabbit import RabbitBroker

from asyncio import Future
from contextlib import asynccontextmanager
from typing import Annotated, Any

@asynccontextmanager
async def lifespan():
    context.set_global("replies_container", {})
    yield

RepliesContainer = Annotated[dict[str, Future[Any]], Context("replies_container")]

broker = RabbitBroker()
app = FastStream(broker, lifespan=lifespan)


@broker.subscriber("in")
@broker.publisher("replies")
async def handle_request(data):
    """Your hardcoded not-FastStream service."""
    print('received task', data)
    return data

@broker.subscriber("in")
async def handle_request(data: dict, cor_id: str = Context("message.correlation_id")):
    """Your hardcoded not-FastStream service."""
    print('收到任务', data, cor_id)
    reply_queue = data.pop('reply_queue', None)
    await broker.publish(message=data, queue=reply_queue, correlation_id=cor_id)