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)
