powerloom · xadahiya · Jun 13, 2023 · Jun 8, 2023 · Jun 8, 2023
diff --git a/pooler/core_api.py b/pooler/core_api.py
@@ -16,6 +16,8 @@
 from pooler.settings.config import settings
 from pooler.utils.data_utils import get_project_epoch_snapshot
 from pooler.utils.data_utils import get_project_finalized_cid
+from pooler.utils.data_utils import get_snapshotter_project_status
+from pooler.utils.data_utils import get_snapshotter_status
 from pooler.utils.default_logger import logger
 from pooler.utils.file_utils import read_json_file
 from pooler.utils.redis.rate_limiter import load_rate_limiter_scripts
@@ -353,3 +355,86 @@ async def get_finalized_cid_for_project_id_epoch_id(
     await incr_success_calls_count(auth_redis_conn, rate_limit_auth_dep)
 
     return data
+
+
+@app.get('/internal/snapshotter/status')
+async def get_snapshotter_overall_status(
+    request: Request,
+    response: Response,
+    rate_limit_auth_dep: RateLimitAuthCheck = Depends(
+        rate_limit_auth_check,
+    ),
+):
+    if not (
+        rate_limit_auth_dep.rate_limit_passed and
+        rate_limit_auth_dep.authorized and
+        rate_limit_auth_dep.owner.active == UserStatusEnum.active
+    ):
+        return inject_rate_limit_fail_response(rate_limit_auth_dep)
+
+    try:
+        snapshotter_status = await get_snapshotter_status(
+            request.app.state.redis_pool,
+        )
+    except Exception as e:
+        rest_logger.exception(
+            'Exception in get_snapshotter_overall_status',
+            e=e,
+        )
+        response.status_code = 500
+        return {
+            'status': 'error',
+            'message': f'Unable to get snapshotter status, error: {e}',
+        }
+
+    auth_redis_conn: aioredis.Redis = request.app.state.auth_aioredis_pool
+    await incr_success_calls_count(auth_redis_conn, rate_limit_auth_dep)
+
+    return snapshotter_status
+
+
+@app.get('/internal/snapshotter/status/{project_id}')
+async def get_snapshotter_project_level_status(
+    request: Request,
+    response: Response,
+    project_id: str,
+    data: bool = False,
+    rate_limit_auth_dep: RateLimitAuthCheck = Depends(
+        rate_limit_auth_check,
+    ),
+):
+    if not (
+        rate_limit_auth_dep.rate_limit_passed and
+        rate_limit_auth_dep.authorized and
+        rate_limit_auth_dep.owner.active == UserStatusEnum.active
+    ):
+        return inject_rate_limit_fail_response(rate_limit_auth_dep)
+
+    try:
+        snapshotter_project_status = await get_snapshotter_project_status(
+            request.app.state.redis_pool,
+            project_id=project_id,
+            with_data=data,
+        )
+    except Exception as e:
+        rest_logger.exception(
+            'Exception in get_snapshotter_project_level_status',
+            e=e,
+        )
+        response.status_code = 500
+        return {
+            'status': 'error',
+            'message': f'Unable to get snapshotter status for project_id: {project_id}, error: {e}',
+        }
+
+    if not snapshotter_project_status:
+        response.status_code = 404
+        return {
+            'status': 'error',
+            'message': f'No snapshotter status found for project_id: {project_id}',
+        }
+
+    auth_redis_conn: aioredis.Redis = request.app.state.auth_aioredis_pool
+    await incr_success_calls_count(auth_redis_conn, rate_limit_auth_dep)
+
+    return snapshotter_project_status.dict(exclude_none=True, exclude_unset=True)
diff --git a/pooler/utils/data_utils.py b/pooler/utils/data_utils.py
@@ -15,8 +15,16 @@
 from pooler.utils.default_logger import logger
 from pooler.utils.file_utils import read_json_file
 from pooler.utils.file_utils import write_json_file
+from pooler.utils.models.data_models import ProjectStatus
+from pooler.utils.models.data_models import SnapshotterIncorrectSnapshotSubmission
+from pooler.utils.models.data_models import SnapshotterMissedSnapshotSubmission
+from pooler.utils.models.data_models import SnapshotterProjectStatus
+from pooler.utils.models.data_models import SnapshotterReportState
+from pooler.utils.models.data_models import SnapshotterStatus
+from pooler.utils.models.data_models import SnapshotterStatusReport
 from pooler.utils.redis.redis_keys import project_finalized_data_zset
 from pooler.utils.redis.redis_keys import project_first_epoch_hmap
+from pooler.utils.redis.redis_keys import project_snapshotter_status_report_key
 from pooler.utils.redis.redis_keys import source_chain_block_time_key
 from pooler.utils.redis.redis_keys import source_chain_epoch_size_key
 from pooler.utils.redis.redis_keys import source_chain_id_key
@@ -347,3 +355,94 @@ async def get_project_epoch_snapshot_bulk(
     )
 
     return all_snapshot_data
+
+
+# get snapshotter high level status for all the projects
+async def get_snapshotter_status(redis_conn: aioredis.Redis):
+    status_keys = []
+
+    all_projects = await redis_conn.smembers('storedProjectIds')
+    all_projects = [project_id.decode('utf-8') for project_id in all_projects]
+
+    for project_id in all_projects:
+        status_keys.append(f'projectID:{project_id}:totalSuccessfulSnapshotCount')
+        status_keys.append(f'projectID:{project_id}:totalIncorrectSnapshotCount')
+        status_keys.append(f'projectID:{project_id}:totalMissedSnapshotCount')
+
+    all_projects_status = await redis_conn.mget(status_keys)
+
+    total_successful_submissions = 0
+    total_incorrect_submissions = 0
+    total_missed_submissions = 0
+    overall_status = SnapshotterStatus(projects=[])
+
+    # project level data
+    project_index = 0
+    successful_submissions = 0
+    incorrect_submissions = 0
+    missed_submissions = 0
+
+    for index, count in enumerate(all_projects_status):
+        if count is None:
+            count = 0
+
+        # as each project has three counts viz. successful, incorrect and missed
+        if index % 3 == 0:
+            successful_submissions = int(count)
+            total_successful_submissions += int(count)
+        elif index % 3 == 1:
+            incorrect_submissions = int(count)
+            total_incorrect_submissions += int(count)
+        else:
+            missed_submissions = int(count)
+            total_missed_submissions += int(count)
+            overall_status.projects.append(
+                ProjectStatus(
+                    projectId=all_projects[project_index],
+                    successfulSubmissions=successful_submissions,
+                    incorrectSubmissions=incorrect_submissions,
+                    missedSubmissions=missed_submissions,
+                ),
+            )
+            project_index += 1
+
+    overall_status.totalSuccessfulSubmissions = total_successful_submissions
+    overall_status.totalIncorrectSubmissions = total_incorrect_submissions
+    overall_status.totalMissedSubmissions = total_missed_submissions
+
+    return overall_status
+
+
+# gets snapshotter status for a particular project
+async def get_snapshotter_project_status(redis_conn: aioredis.Redis, project_id: str, with_data: bool):
+    reports = await redis_conn.hgetall(project_snapshotter_status_report_key(project_id))
+
+    reports = {
+        int(k.decode('utf-8')): SnapshotterStatusReport.parse_raw(v.decode('utf-8'))
+        for k, v in reports.items()
+    }
+
+    project_status = SnapshotterProjectStatus(missedSubmissions=[], incorrectSubmissions=[])
+
+    for epoch_id, report in reports.items():
+        if report.state is SnapshotterReportState.MISSED_SNAPSHOT:
+            project_status.missedSubmissions.append(
+                SnapshotterMissedSnapshotSubmission(
+                    epochId=epoch_id,
+                    reason=report.reason,
+                    finalizedSnapshotCid=report.finalizedSnapshotCid,
+                ),
+            )
+        elif report.state is SnapshotterReportState.SUBMITTED_INCORRECT_SNAPSHOT:
+            project_status.incorrectSubmissions.append(
+                SnapshotterIncorrectSnapshotSubmission(
+                    epochId=epoch_id,
+                    reason=report.reason,
+                    submittedSnapshotCid=report.submittedSnapshotCid,
+                    submittedSnapshot=report.submittedSnapshot if with_data else None,
+                    finalizedSnapshotCid=report.finalizedSnapshotCid,
+                    finalizedSnapshot=report.finalizedSnapshot if with_data else None,
+                ),
+            )
+
+    return project_status
diff --git a/pooler/utils/models/data_models.py b/pooler/utils/models/data_models.py
@@ -1,4 +1,7 @@
+from enum import Enum
+from typing import Any
 from typing import Dict
+from typing import List
 from typing import Optional
 
 from pydantic import BaseModel
@@ -69,3 +72,63 @@ class ProjectSpecificState(BaseModel):
 class ProtocolState(BaseModel):
     project_specific_states: Dict[str, ProjectSpecificState]  # project ID -> project specific state
     synced_till_epoch_id: int
+
+
+class SnapshotterReportState(Enum):
+    MISSED_SNAPSHOT = 'MISSED_SNAPSHOT'
+    SUBMITTED_INCORRECT_SNAPSHOT = 'SUBMITTED_INCORRECT_SNAPSHOT'
+
+
+class ProjectStatus(BaseModel):
+    projectId: str
+    successfulSubmissions: int = 0
+    incorrectSubmissions: int = 0
+    missedSubmissions: int = 0
+
+
+class SnapshotterStatus(BaseModel):
+    totalSuccessfulSubmissions: int = 0
+    totalIncorrectSubmissions: int = 0
+    totalMissedSubmissions: int = 0
+    projects: List[ProjectStatus]
+
+
+class SnapshotterMissedSubmission(BaseModel):
+    epochId: int
+    reason: str
+
+
+class SnapshotterIncorrectSubmission(BaseModel):
+    epochId: int
+    incorrectCid: str
+    payloadDump: str
+    reason: str = ''
+
+
+class SnapshotterStatusReport(BaseModel):
+    submittedSnapshotCid: str
+    submittedSnapshot: Dict[str, Any] = {}
+    finalizedSnapshotCid: str
+    finalizedSnapshot: Dict[str, Any] = {}
+    state: SnapshotterReportState
+    reason: str = ''
+
+
+class SnapshotterMissedSnapshotSubmission(BaseModel):
+    epochId: int
+    finalizedSnapshotCid: str
+    reason: str
+
+
+class SnapshotterIncorrectSnapshotSubmission(BaseModel):
+    epochId: int
+    submittedSnapshotCid: str
+    submittedSnapshot: Optional[Dict[str, Any]]
+    finalizedSnapshotCid: str
+    finalizedSnapshot: Optional[Dict[str, Any]]
+    reason: str = ''
+
+
+class SnapshotterProjectStatus(BaseModel):
+    missedSubmissions: List[SnapshotterMissedSnapshotSubmission]
+    incorrectSubmissions: List[SnapshotterIncorrectSnapshotSubmission]
diff --git a/pooler/utils/redis/redis_keys.py b/pooler/utils/redis/redis_keys.py
@@ -55,6 +55,7 @@
 def project_finalized_data_zset(project_id):
     return f'projectID:{project_id}:finalizedData'
 
+
 # project first epoch hashmap
 
 
@@ -76,3 +77,23 @@ def source_chain_epoch_size_key():
 
 def project_last_finalized_epoch_key(project_id):
     return f'projectID:{project_id}:lastFinalizedEpoch'
+
+
+def project_successful_snapshot_submissions_suffix():
+    return 'totalSuccessfulSnapshotCount'
+
+
+def project_incorrect_snapshot_submissions_suffix():
+    return 'totalIncorrectSnapshotCount'
+
+
+def project_missed_snapshot_submissions_suffix():
+    return 'totalMissedSnapshotCount'
+
+
+def project_snapshotter_status_report_key(project_id):
+    return f'projectID:{project_id}:snapshotterStatusReport'
+
+
+def stored_projects_key():
+    return 'storedProjectIds'