Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add pause command to send bot to idle state #48

Merged
merged 6 commits into from
Oct 27, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 12 additions & 9 deletions src/api/rate_limiter.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,17 +58,20 @@ def __init__(self, request_timeout=None, timeout_delay=None, idle_state_sleep=No
else:
self._max_attempt = DEFAULT_MAX_ATTEMPT

self._pause = False
self._waiting = False
self._idle = False
self._request_timeout = request_timeout or DEFAULT_REQUEST_TIMEOUT
self._timeout_delay = timeout_delay or DEFAULT_TIMEOUT_DELAY
self._idle_state_sleep = idle_state_sleep or DEFAULT_IDLE_STATE_SLEEP
self.task = asyncio.create_task(self.handle_requests())
self.logger = logging.getLogger(__name__)

def go_idle(self):
def go_idle(self, manually = False):
self._idle = True
self.logger.error("Entering idle state due to previous errors")
if manually:
self.logger.error("Entering idle state due to manual command")
else:
self.logger.error("Entering idle state due to previous errors")

def is_idle(self):
return self._idle
Expand All @@ -77,13 +80,13 @@ def exit_idle(self):
self.logger.info("Resuming from idle state")
self._idle = False

def is_paused(self):
return self._pause
def is_waiting(self):
return self._waiting

async def pause(self, delay):
self._pause = True
async def wait(self, delay):
self._waiting = True
await asyncio.sleep(delay)
self._pause = False
self._waiting = False

def handle_get_request(self, request):
try:
Expand Down Expand Up @@ -162,7 +165,7 @@ async def handle_requests(self):
self.requests[request.key]["result"] = self.handle_get_request(request)
except RateLimiterError as exc:
if isinstance(exc, RLErrorWithPause):
await self.pause(exc.time_to_wait)
await self.wait(exc.time_to_wait)

if request.attempt < self._max_attempt:
await self.queue.put(request)
Expand Down
3 changes: 3 additions & 0 deletions src/api/rootme_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ def __init__(self, rate_limiter: RateLimiter):
raise RuntimeError("API_URL is not set.")
self.rate_limiter = rate_limiter

def get_rate_limiter(self):
return self.rate_limiter

async def get_challenge_by_id(self, _id):
"""
Get a challenge from the API
Expand Down
21 changes: 18 additions & 3 deletions src/bot/root_pythia_cogs.py
Original file line number Diff line number Diff line change
Expand Up @@ -36,25 +36,40 @@ async def status(self, ctx):
check = ":white_check_mark:"
cross = ":x:"
rl_alive = cross if rate_limiter.task.done() else check
rl_paused = check if rate_limiter.is_paused() else cross
rl_waiting = check if rate_limiter.is_waiting() else cross
rl_idle = check if rate_limiter.is_idle() else cross
# pylint: disable-next=no-member
bot_loop_alive = check if self.check_new_solves.is_running() else cross
await ctx.send(
f"RootMe API's Rate Limiter status:\n"
f"- alive: {rl_alive}\n"
f" - paused: {rl_paused}\n"
f" - waiting: {rl_waiting}\n"
f" - idle: {rl_idle}\n"
f"Bot's `check_new_solves` loop:\n"
f"- alive: {bot_loop_alive}\n"
)

@commands.command(name="pause")
async def pause(self, ctx):
"""
The bot enters its idle state
"""
rate_limiter = self.dbmanager.get_api_manager().get_rate_limiter()
if rate_limiter.is_idle():
await ctx.message.channel.send("The Rate Limiter is already idle, can't pause.")
return

rate_limiter.go_idle(manually = True)
await ctx.message.channel.send(
"Bot is in idle state, no requests will be sent."
)

@commands.command(name="resume")
async def resume(self, ctx):
"""
The bot leaves its idle state
"""
rate_limiter = self.dbmanager.api_manager.rate_limiter
rate_limiter = self.dbmanager.get_api_manager().get_rate_limiter()
if not rate_limiter.is_idle():
await ctx.message.channel.send("The Rate Limiter isn't idle, no need to resume.")
return
Expand Down
3 changes: 3 additions & 0 deletions src/database/db_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,9 @@ def _init_db(self):
cur.execute(sql_create_user_table)
cur.close()

def get_api_manager(self):
return self.api_manager

async def add_user(self, idx):
"""Call the API Manager to get a user by his id then create a User object and store it"""
cur = self.db.cursor()
Expand Down
Loading