class Manager(manager_grpc.ManagerServicer):
"""
Implements a server for the Manager gRPC protocol.
"""
def __init__(self, total_shards):
"""
Instantiates a new manager server that handles some
number of shards.
"""
logger.info(f"Number of shards: {total_shards}")
self.hoarfrost_gen = HoarFrostGenerator()
self.total_shards = total_shards
self.registered = [False for _ in range(total_shards)]
self.last_checkin = dict()
self.store = dict()
def health_check(self):
while True:
time.sleep(5)
for shard, last_checkin in self.last_checkin.items():
if last_checkin is not None and last_checkin < datetime.now(
) - timedelta(seconds=5):
logger.error(
f"--- SHARD {shard} MISSED ITS HEARTBEAT, DEREGISTERING... ---"
)
self.registered[shard] = False
self.last_checkin[shard] = None
def register(self, request, context):
"""Returns the next shard id that needs to be filled as well as the total shards"""
if all(self.registered):
raise Exception("Shard trying to register even though we're full")
i = next(i for i in range(self.total_shards) if not self.registered[i])
logger.info(
f"Shard requested id, assigning {i + 1}/{self.total_shards}...")
self.registered[i] = True
return message.ShardInfo(shard_id=i, shard_count=self.total_shards)
def guild_count(self, request, context):
"""Return guild and user count information"""
gc = 0
uc = 0
for guilds in self.store.values():
gc += len(guilds)
for guild in guilds:
uc += guild.member_count
return message.GuildInfo(guild_count=gc, user_count=uc)
def checkin(self, request, context):
self.last_checkin[request.shard_id] = datetime.now()
self.registered[request.shard_id] = True
return message.CheckInResponse()
def publish_file(self, request_iterator, context):
"""Missing associated documentation comment in .proto file"""
first = next(request_iterator)
filetype = "png" if first.filetype == "" else first.filetype
name = first.name
if name == "":
name = str(self.hoarfrost_gen.generate())
location = first.location
if location == "":
location = "assets"
directory = f"/var/www/{location}"
if not os.path.exists(directory):
os.makedirs(directory)
with open(f"{directory}/{name}.{filetype}", "wb") as f:
logger.info(f"Writing {directory}/{name}.{filetype}")
f.write(first.file)
for datum in request_iterator:
f.write(datum.file)
return message.Url(
url=f"https://cdn.{domain_name}/{location}/{name}.{filetype}")
def all_guilds(self, request, context):
"""Return information about all guilds that the bot is in, including their admins"""
for guilds in self.store.values():
for guild in guilds:
yield guild
def guild_update(self, request_iterator, context):
"""Update the manager with the latest information about a shard's guilds"""
guilds = []
for guild in request_iterator:
guilds.append(guild)
if len(guilds) == 0:
return message.UpdateResponse()
logger.debug(
f"Received guild list from shard {guilds[0].shard_id + 1} of {len(guilds)} guilds"
)
self.store[guilds[0].shard_id] = guilds
return message.UpdateResponse()
class Manager:
"""Manages shards. assigns shard nodes their ids and checks if they are alive"""
def __init__(self, total_shards):
logger.info(f"Number of shards: {total_shards}")
self.hoarfrost_gen = HoarFrostGenerator()
self.total_shards = total_shards
self.registered = [False for _ in range(total_shards)]
self.last_checkin = {}
self.store = {}
async def handle_task(self, method, *args, **kwargs):
try:
return (await (getattr(self, method)(*args, **kwargs)), 200)
except Exception as e:
logger.exception(f"caught: '{e}' while executing '{method}'")
return {'message': f"caught: '{e}' while executing '{method}'"}
async def health_check(self):
while True:
await asyncio.sleep(1)
for shard, last_checkin in self.last_checkin.items():
if last_checkin is not None and last_checkin < datetime.now() - timedelta(seconds=5):
logger.error(f"--- SHARD {shard} MISSED ITS HEARTBEAT, DEREGISTERING... ---")
self.registered[shard] = False
self.last_checkin[shard] = None
async def register(self):
"""Returns the next shard id that needs to be filled as well as the total shards"""
if all(self.registered):
raise Exception("Shard trying to register even though we're full")
i = next(i for i in range(self.total_shards) if not self.registered[i])
logger.info(f'Shard requested id, assigning {i + 1}/{self.total_shards}...')
self.store[i] = {'shard_id': i}
self.registered[i] = True
return {'shard_id': i, 'shard_count': self.total_shards}
async def all_guilds(self):
"""Return information about all guilds that the bot is in, including their admins"""
guilds = []
for shard, shard_store in self.store.items():
guilds += shard_store.get('guilds', ())
return guilds
async def guild_count(self):
"""Return guild and user count information"""
guild_count = 0
user_count = 0
for shard, shard_store in self.store.items():
guilds = shard_store.get('guilds', ())
guild_count += len(guilds)
for guild in guilds:
user_count += guild['member_count']
return {'guild_count': guild_count, 'user_count': user_count}
async def guild_update(self, shard_id, guilds):
"""Update the manager with the latest information about a shard's guilds"""
logger.debug(f"someone sent guild list containing {len(guilds)} guilds")
self.store[int(shard_id)]['guilds'] = guilds
return {"message": "thanks"}
async def checkin(self, shard_id):
self.last_checkin[shard_id] = datetime.now()
self.registered[shard_id] = True
return "nice"
async def publish_file(self, location: str = 'assets', name: str = '', filetype: str = 'png', data: str = ''):
assert data != ''
if name == '':
name = str(self.hoarfrost_gen.generate())
directory = f'/var/www/{location}'
if not os.path.exists(directory):
os.makedirs(directory)
with open(f'{directory}/{name}.{filetype}', 'wb') as f:
logger.info(f'writing {directory}/{name}.{filetype}')
f.write(base64.b64decode(data))
return {'url': f'https://cdn.{domain_name}/{location}/{name}.{filetype}'}
def generate_response_ast(response):
response = generate_response(response)
try:
return json.dumps(parse(response).stringify())
except Exception as e:
print(e.__class__.__name__)
print(f"caught while parsing: '{response}'")
return ""
responses = defaultdict(list)
for cmd in command_list:
trigger = generate_trigger(cmd.trigger, cmd.server_id)
responses[cmd.server_id].append(AutoResponse(
hoarfrost_gen.generate(),
trigger,
generate_response(cmd.response),
cmd.author_id,
cmd.server_id,
generate_trigger_regex(trigger, get_mode(trigger)),
get_puncutation(trigger),
generate_response_ast(cmd.response),
get_mode(trigger),
cmd.count
))
for guild_id in responses.keys():
reggys = []
uses_regex = False
for resp in responses[guild_id]: