async def get_stub_for_udid(self, udid: Optional[str]) -> CompanionClient: is_companion_available = ( self.is_companion_available_for_target_udid(udid) if udid else False) if udid and udid in self._stub_map and udid in self._udid_companion_map: return CompanionClient( stub=self._stub_map[udid], is_local=self._udid_companion_map[udid].is_local, udid=udid, logger=self._logger, is_companion_available=is_companion_available, ) else: async with self.create_companion_for_target_with_udid( target_udid=udid) as companion: stub = self.get_stub_for_address(companion.host, none_throws(companion.port)) self._stub_map[companion.udid] = stub return CompanionClient( stub=stub, is_local=companion.is_local, udid=udid, logger=self._logger, is_companion_available=is_companion_available, )
async def _tramp(stream: Stream[_T, _U], *args: Any, **kwargs: Any) -> None: partial_call = call if _takes_client(call): client = await companion_provider(stream.metadata.get("udid")) logger = client.logger.getChild(name) companion_client = CompanionClient( stub=client.stub, is_local=client.is_local, udid=client.udid, logger=logger, is_companion_available=client.is_companion_available, ) partial_call = partial(partial_call, client=companion_client) if _takes_context(call): context = await context_provider() partial_call = partial(partial_call, context=context) if _takes_stream(call): # pyre-fixme[29]: `Union[Callable[..., Any], Callable[[*(Any), **(Any)], # Any]]` is not a function. await partial_call(stream=stream) else: request = await stream.recv_message() # pyre-fixme[29]: `Union[Callable[..., Any], Callable[[*(Any), **(Any)], # Any]]` is not a function. response = await partial_call(request=request) await stream.send_message(response)
async def _make_client() -> CompanionClient: client = await daemon_provider() return CompanionClient( stub=MetadataStubInjector( stub=client.stub, metadata={"udid": client.udid} if client.udid else {}), is_local=client.is_local, udid=client.udid, logger=client.logger.getChild(name), )
async def provide_client(self) -> CompanionClient: await self.daemon_spawner.start_daemon_if_needed( force_kill=self.force_kill_daemon) if not self.channel or not self.stub: self.channel = Channel(self.host, self.port, loop=asyncio.get_event_loop()) self.stub = CompanionServiceStub(channel=self.channel) return CompanionClient(stub=self.stub, is_local=True, udid=self.target_udid, logger=self.logger)