It seems like I have a memory leak problem in my FastAPI app. It has only one endpoint that works with Websockets. I will give you the code for that endpoint, use case, profile service and connection manager. Find the memory leak and help to identify the exact place to fix it.

Router:
@profile_router.websocket("/ws/{connection_state}")
async def websocket_endpoint(
    websocket: WebSocket,
    connection_state: str,
    usecase: BaseUseCase = UseCase(HandleWebsocketConnectionUseCase),
) -> None:
    await usecase.execute(websocket, connection_state)

Usecase:

class HandleWebsocketConnectionUseCase(BaseUseCase):
    def __init__(
        self,
        token_service: TokenServiceInterface,
        profile_service: ProfileServiceInterface,
        rate_limiter: RateLimiter,
    ) -> None:
        self._token_service = token_service
        self._profile_service = profile_service
        self._rate_limiter = rate_limiter

    async def execute(self, websocket: WebSocket, connection_state: str) -> None:
        try:
            await self._profile_service.accept_ws_connection(
                websocket, connection_state
            )
            auth_message = await websocket.receive_text()
            token, ref_code = await self._profile_service.get_auth_token_and_ref_code(
                auth_message
            )
        except AuthenticationError:
            raise NotAllowedException
        except (WebSocketDisconnect, RuntimeError):
            return
        try:
            token_data = await self._token_service.get_access_token_data(token)
        except ExpiredTokenError:
            raise TokenExpiredException
        except (InvalidTokenError, NoTokenError):
            raise NotAllowedException

        await self._profile_service.connect_websocket(
            auth_user_uuid=token_data.auth_user_uuid,
            state=connection_state,
            websocket=websocket,
        )
        try:
            await self._profile_service.authorize_user(
                auth_user_uuid=token_data.auth_user_uuid,
                name=token_data.name,
                state=connection_state,
                ref_code=ref_code,
            )
        except AccountDeactivatedError:
            await websocket.close(code=4023, reason=ErrorCode.PROFILE_DEACTIVATED)
            await self._profile_service.disconnect_websocket(
                auth_user_uuid=token_data.auth_user_uuid, state=connection_state
            )
            return

        await self._rate_limiter.start_timeout(connection_state, websocket)
        try:
            while True:
                command_message = await websocket.receive_text()

                if command_message == "ping":
                    await websocket.send_text("pong")
                    continue

                await self._rate_limiter.reset_timeout(connection_state, websocket)

                if await self._rate_limiter.is_rate_limited(token_data.auth_user_uuid):
                    raise RateLimitExceededException

                await self._profile_service.handle_command(
                    auth_user_uuid=token_data.auth_user_uuid,
                    telegram_id=token_data.telegram_id,
                    state=connection_state,
                    command_message=command_message,
                )
        except AccountDeactivatedError:
            await websocket.close(code=4023, reason=ErrorCode.PROFILE_DEACTIVATED)
            await self._rate_limiter.end_timeout(connection_state)
            await self._profile_service.disconnect_websocket(
                auth_user_uuid=token_data.auth_user_uuid, state=connection_state
            )
        except (WebSocketDisconnect, RuntimeError):
            await self._rate_limiter.end_timeout(connection_state)
            await self._profile_service.disconnect_websocket(
                auth_user_uuid=token_data.auth_user_uuid, state=connection_state
            )
        finally:
            await self._rate_limiter.end_timeout(connection_state)
            await self._profile_service.disconnect_websocket(
                auth_user_uuid=token_data.auth_user_uuid, state=connection_state
            )

Profile service:

class ProfileService(ProfileServiceInterface):
    def __call__(self):
        return self

    def __init__(
        self,
        redis: RedisManager,
        pulsar: PulsarManager,
        ws: ConnectionManager,
        action_repository: ActionRepository,
        profile_repository: ProfileRepository,
        task_repository: TaskRepository,
        command_dispatcher: CommandDispatcher,
    ) -> None:
        self._redis_manager = redis
        self._pulsar_manager = pulsar
        self._ws_manager = ws
        self._pulsar_manager.set_updates_consumer_with_callback(self._handle_pubsub)
        self._action_repository = action_repository
        self._profile_repository = profile_repository
        self._task_repository = task_repository
        self._command_dispatcher = command_dispatcher

    def _handle_pubsub(self, consumer, message):
        event = self._pulsar_manager.consume_update(message)
        updates = ProfileResponse(
            profile_uuid=event.profile_uuid,
            auth_user_uuid=event.auth_user_uuid,
            is_ambassador=event.is_ambassador,
            ref_code=event.ref_code,
            ref_quantity=event.ref_quantity,
            name=event.name,
            direct_points=event.direct_points,
            ref_points=event.ref_points,
            crystal_balls=event.crystal_balls,
            streak=event.streak,
            actions=event.actions,
        )
        try:
            loop = asyncio.get_running_loop()
            loop.call_soon_threadsafe(
                asyncio.create_task,
                self._ws_manager.handle_pubsub_message(updates=updates),
            )
        except RuntimeError:
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            loop.run_until_complete(
                self._ws_manager.handle_pubsub_message(updates=updates)
            )

    @staticmethod
    async def accept_ws_connection(websocket: WebSocket, connection_state: str) -> None:
        await websocket.accept()
        if not validate_state(connection_state):
            raise AuthenticationError

    async def get_auth_token_and_ref_code(
        self, auth_message: str
    ) -> Tuple[str, Optional[str]]:
        try:
            auth_message_data = AuthCommand.model_validate_json(auth_message).payload
            return auth_message_data.token, auth_message_data.ref_code
        except (ValueError, KeyError):
            raise AuthenticationError

    async def connect_websocket(
        self, auth_user_uuid: uuid.UUID, state: str, websocket: WebSocket
    ) -> None:
        await self._ws_manager.connect(
            auth_user_uuid=auth_user_uuid, state=state, websocket=websocket
        )

    async def authorize_user(
        self,
        auth_user_uuid: uuid.UUID,
        name: Optional[str],
        state: str,
        ref_code: Optional[str],
    ) -> None:
        profile_data = await self._get_user_profile(auth_user_uuid)
        if profile_data:
            response = Response(
                type=ResponseType.PROFILE, response=profile_data
            ).model_dump_json()
            await self._ws_manager.send_personal_message(
                auth_user_uuid, state, response
            )
            return

        await self._create_new_profile(auth_user_uuid, name, ref_code)

    async def _create_new_profile(
        self, auth_user_uuid: uuid.UUID, name: Optional[str], ref_code: Optional[str]
    ) -> None:
        direct_referrer_uuid, indirect_referrer_uuid = (
            await self._profile_repository.get_referrer_uuids(ref_code=ref_code)
        )
        events = [
            PulsarCommandEvent(
                payload=PulsarCreateProfileEvent(
                    auth_user_uuid=auth_user_uuid,
                    ref_code=generate_random_alphanum(8).upper(),
                    name=name or f"MiniappUser-{generate_random_alphanum(6)}",
                    direct_referrer_uuid=direct_referrer_uuid,
                    indirect_referrer_uuid=indirect_referrer_uuid,
                )
            )
        ]
        if direct_referrer_uuid:
            events.append(
                PulsarCommandEvent(
                    payload=PulsarClaimEvent(
                        action=ActionType.REFERRAL_ACQUIRE,
                        value=1,
                        auth_user_uuid=direct_referrer_uuid,
                        idempotency_key=auth_user_uuid,
                    )
                )
            )
        await self._pulsar_manager.produce_commands(events)

    async def _get_user_profile(
        self, auth_user_uuid: uuid.UUID
    ) -> Optional[ProfileResponse]:
        profile_data = await self._redis_manager.get_cached_profile(auth_user_uuid)
        if profile_data:
            actions = {k: v.updated_at for k, v in profile_data.actions.items()}
        else:
            profile_data = await self._profile_repository.get_by_auth_user_uuid(
                auth_user_uuid
            )
            if not profile_data:
                return None

            actions = await self._action_repository.get_actions_by_profile_uuid(
                profile_data.uuid
            )
            await self._redis_manager.set_cached_profile(profile_data, actions)

            actions = {ac.action: ac.updated_at for ac in actions}

        if profile_data.is_deactivated:
            raise AccountDeactivatedError

        streak = await self._redis_manager.get_streak(auth_user_uuid)

        return ProfileResponse.model_validate(
            {
                "profile_uuid": profile_data.uuid,
                **profile_data.model_dump(),
                "streak": streak,
                "actions": actions,
            }
        )

    async def _get_global_leaderboard(self) -> GlobalLeaderboardResponse:
        leaderboard_data = await self._redis_manager.get_leaderboard()
        if not leaderboard_data:
            async with self._redis_manager.lock(LEADERBOARD_LOCK):
                leaderboard_data = await self._redis_manager.get_leaderboard()
                if not leaderboard_data:
                    leaderboard = (
                        await self._profile_repository.get_global_leaderboard()
                    )
                    leaderboard_data = await self._redis_manager.set_leaderboard(
                        leaderboard
                    )

        return GlobalLeaderboardResponse.model_validate(leaderboard_data.model_dump())

    async def _get_total_users_quantity(self) -> int:
        total_users = await self._redis_manager.get_total_users()
        if not total_users:
            async with self._redis_manager.lock(TOTAL_USERS_QUANTITY_LOCK):
                total_users = await self._redis_manager.get_total_users()
                if not total_users:
                    total_users = (
                        await self._profile_repository.get_total_users_quantity()
                    )
                    await self._redis_manager.set_total_users(total_users)
        return total_users

    async def _get_user_rank(self, auth_user_uuid: uuid.UUID) -> int:
        user_rank = await self._redis_manager.get_rank(auth_user_uuid)
        if not user_rank:
            lock = f"{RANK_LOCK}:{auth_user_uuid}"
            async with self._redis_manager.lock(lock):
                user_rank = await self._redis_manager.get_rank(auth_user_uuid)
                if not user_rank:
                    user_rank = await self._profile_repository.get_leaderboard_rank(
                        auth_user_uuid
                    )
                    await self._redis_manager.set_rank(auth_user_uuid, user_rank)
        return user_rank

    async def _get_leaderboard_and_rank(
        self, auth_user_uuid: uuid.UUID
    ) -> Tuple[GlobalLeaderboardResponse, int]:
        leaderboard = await self._get_global_leaderboard()
        for idx, entry in enumerate(leaderboard.leaderboard):
            if entry.auth_user_uuid == auth_user_uuid:
                user_rank = idx + 1
                return leaderboard, user_rank

        user_rank = await self._get_user_rank(auth_user_uuid)
        return leaderboard, user_rank

    async def _get_user_referrals(
        self, auth_user_uuid: uuid.UUID, ref_quantity: int
    ) -> ReferralsResponse:
        if ref_quantity == 0:
            return ReferralsResponse(referrals=[])
        referrals_data = await self._redis_manager.get_referrals(auth_user_uuid)
        if not referrals_data:
            lock = f"{REFERRALS_LOCK}:{auth_user_uuid}"
            async with self._redis_manager.lock(lock):
                referrals_data = await self._redis_manager.get_referrals(auth_user_uuid)
                if not referrals_data:
                    referrals = await self._profile_repository.get_referrals(
                        auth_user_uuid
                    )
                    referrals_data = await self._redis_manager.set_referrals(
                        auth_user_uuid, referrals
                    )
        return ReferralsResponse.model_validate(referrals_data.model_dump())

    async def _get_tasks(self) -> TasksValue:
        tasks_data = await self._redis_manager.get_tasks()
        if not tasks_data:
            async with self._redis_manager.lock(TASKS_LOCK):
                tasks_data = await self._redis_manager.get_tasks()
                if not tasks_data:
                    tasks = await self._task_repository.get_all_tasks()
                    tasks_data = await self._redis_manager.set_tasks(tasks)

        return tasks_data

    async def disconnect_websocket(
        self, auth_user_uuid: uuid.UUID, state: Optional[str]
    ):
        await self._ws_manager.disconnect(auth_user_uuid=auth_user_uuid, state=state)

    async def handle_command(
        self,
        auth_user_uuid: uuid.UUID,
        telegram_id: Optional[str],
        state: str,
        command_message: str,
    ) -> None:
        try:
            command_handler = await self._command_dispatcher.dispatch(command_message)
            command_context = await self._get_command_context(
                auth_user_uuid=auth_user_uuid,
                telegram_id=telegram_id,
                command=command_handler.command,
            )
            command_events, command_notification = await command_handler.handle(
                command_context
            )
            await self._pulsar_manager.produce_commands(command_events)
            if command_notification:
                await self._ws_manager.send_personal_message(
                    auth_user_uuid=auth_user_uuid,
                    state=state,
                    message=command_notification.model_dump_json(),
                )
        except ServiceError as exc:
            await self._handle_error(
                auth_user_uuid=auth_user_uuid,
                state=state,
                error_message=exc.code,
                data=exc.data,
            )
            return

    async def _get_command_context(
        self,
        auth_user_uuid: uuid.UUID,
        telegram_id: Optional[str],
        command: CommandType,
    ) -> CommandExecutionContext:
        rank = None
        leaderboard = None
        total_users = None
        referrals = None
        tasks = None

        profile_data = await self._redis_manager.get_cached_profile(auth_user_uuid)
        if profile_data is None:
            raise ProfileNotFoundError
        streak = await self._redis_manager.get_streak(auth_user_uuid)
        base_value = COMMAND_TO_BASE_VALUE_MAPPING.get(command, None)

        if command == CommandType.FETCH_LEADERBOARD:
            leaderboard, rank = await self._get_leaderboard_and_rank(auth_user_uuid)
            total_users = await self._get_total_users_quantity()
        if command == CommandType.FETCH_REFERRALS:
            referrals = await self._get_user_referrals(
                auth_user_uuid, profile_data.ref_quantity
            )
        if command in (
            CommandType.FETCH_TASKS,
            CommandType.CLAIM_TASK_POINTS,
            CommandType.VERIFY_TASK,
        ):
            tasks = await self._get_tasks()

        return CommandExecutionContext(
            profile=profile_data,
            telegram_id=telegram_id,
            streak=streak,
            base_value=base_value,
            rank=rank,
            leaderboard=leaderboard,
            referrals=referrals,
            total_users=total_users,
            tasks_data=tasks,
        )

    async def _handle_error(
        self,
        auth_user_uuid: uuid.UUID,
        state: str,
        error_message: str,
        data: Optional[Dict[str, Any]] = None,
    ):
        await self._ws_manager.send_error_message(
            auth_user_uuid=auth_user_uuid,
            state=state,
            error_message=error_message,
            data=data,
        )

Connection manager:

class ConnectionManager:
    def __init__(self) -> None:
        self.active_connections: Dict[uuid.UUID, Dict[str, WebSocket]] = {}
        self.lock: asyncio.Lock = asyncio.Lock()

    async def connect(
        self, auth_user_uuid: uuid.UUID, state: str, websocket: WebSocket
    ):
        async with self.lock:
            current_user_connections = self.active_connections.get(auth_user_uuid, {})

            if len(current_user_connections) >= MAX_CONNECTIONS_PER_USER:
                await websocket.close(code=4002, reason=ErrorCode.TOO_MANY_CONNECTIONS)
                return

            current_user_connections.update({state: websocket})
            self.active_connections[auth_user_uuid] = current_user_connections

    async def disconnect(self, auth_user_uuid: uuid.UUID, state: str) -> None:
        current_user_connections = self.active_connections.get(auth_user_uuid, None)
        if current_user_connections:
            websocket = current_user_connections.get(state, None)
            if (
                websocket
                and websocket.application_state is not WebSocketState.DISCONNECTED
            ):
                await websocket.close()

            async with self.lock:
                current_user_connections.pop(state, None)
                self.active_connections[auth_user_uuid] = current_user_connections
        if not current_user_connections:
            self.active_connections.pop(auth_user_uuid, None)

    async def send_message(self, auth_user_uuid: uuid.UUID, message: str) -> None:
        current_user_connections = self.active_connections.get(auth_user_uuid, {})

        for websocket in current_user_connections.values():
            if websocket:
                await websocket.send_text(message)

    async def send_personal_message(
        self, auth_user_uuid: uuid.UUID, state: str, message: str
    ) -> None:
        websocket = self.active_connections.get(auth_user_uuid, {}).get(state, None)
        if websocket:
            await websocket.send_text(message)

    async def send_error_message(
        self,
        auth_user_uuid: uuid.UUID,
        state: str,
        error_message: str,
        data: Optional[Dict[str, Any]] = None,
    ):
        websocket = self.active_connections.get(auth_user_uuid, {}).get(state, None)
        if websocket:
            await websocket.send_text(
                Response(
                    type=(
                        ResponseType.ERROR
                        if data is None
                        else ResponseType.ERROR_DETAILED
                    ),
                    response=(
                        ErrorMessageResponse(message=error_message)
                        if data is None
                        else ErrorDataResponse(
                            code=data.pop("code"), message=error_message, data=data
                        )
                    ),
                ).model_dump_json()
            )

    async def handle_pubsub_message(self, updates: ProfileResponse) -> None:
        response = Response(
            type=ResponseType.PROFILE, response=updates
        ).model_dump_json()
        await self.send_message(auth_user_uuid=updates.auth_user_uuid, message=response)