Allow checking the latency to specified homeservers

This commit is contained in:
Nexus 2024-08-19 18:57:10 +01:00
parent dd1dc066fc
commit 54af79e84d

View file

@ -1,3 +1,5 @@
import urllib.parse
import niobot import niobot
import time import time
import httpx import httpx
@ -5,21 +7,26 @@ import httpx
class LatencyModule(niobot.Module): class LatencyModule(niobot.Module):
@niobot.command("latency") @niobot.command("latency")
async def latency(self, ctx: niobot.Context): async def latency(self, ctx: niobot.Context, homeserver: str = None):
"""See the bot's latency.""" """See the bot's latency."""
latency = ctx.latency latency = ctx.latency
homeserver = await niobot.resolve_homeserver(ctx.message.sender.split(":")[1]) homeserver = homeserver or await niobot.resolve_homeserver(ctx.message.sender.split(":")[1])
if not homeserver.startswith("http"):
homeserver = f"https://{homeserver}"
parsed = urllib.parse.urlparse(homeserver)
loc = "https://" + parsed.netloc
async with httpx.AsyncClient(headers={"User-Agent": niobot.__user_agent__}) as client: async with httpx.AsyncClient(headers={"User-Agent": niobot.__user_agent__}) as client:
timings = [] timings = []
for rnd in range(5): for rnd in range(5):
start = time.time() start = time.time()
await client.get(f"{homeserver}/_matrix/client/v1/ping") await client.get(f"{loc}/_matrix/client/v1/ping")
end = time.time() end = time.time()
timings.append(end - start) timings.append(end - start)
fed_latency = sum(timings) / len(timings) * 1000 fed_latency = sum(timings) / len(timings) * 1000
hs_target = ctx.message.sender.split(":")[1] hs_target = parsed.hostname
return await ctx.respond( return await ctx.respond(
"Latency: {:,.2f}ms (federation latency to {!r}: {:,.2f}ms)".format(latency, hs_target, fed_latency) "Latency: {:,.2f}ms (federation latency to {!r}: {:,.2f}ms)".format(latency, hs_target, fed_latency)
) )