Prevent automatic shutdown in Docker/Linux (#119)

* Manage a single Scheduler instance and inject it as a dependency

* Switch the docker image over to the official gunicorn/uvicorn image

* Fix linter errors

Author: rkorsak

Dockerfile

@@ -1,6 +1,5 @@
-FROM python:3.8 AS base
-ENV host 0.0.0.0
-ENV port 8000
+FROM tiangolo/uvicorn-gunicorn-fastapi:python3.8 AS base
+ENV PORT 8000
 RUN apt update && apt-get install -y \
     libgmp-dev \
     libmpfr-dev \
@@ -10,14 +9,12 @@ COPY ./pyproject.toml /tmp/
 COPY ./poetry.lock /tmp/
 RUN cd /tmp && poetry export -f requirements.txt > requirements.txt
 RUN pip install -r /tmp/requirements.txt
+EXPOSE $PORT
 
 FROM base AS dev
-VOLUME [ "/app" ]
-EXPOSE $port
-CMD uvicorn app.main:app --reload --host "$host" --port "$port"
+VOLUME [ "/app/app" ]
+CMD /start-reload.sh
 
 FROM base AS prod
-COPY ./app /app
-EXPOSE $port
-# TODO: We should not have to use the --reload flag here! See issue #80
-CMD uvicorn app.main:app --reload --host "$host" --port "$port"
+COPY ./app /app/app
+# The base image will start gunicorn

app/api/v1/guardian/ballot.py

@@ -4,8 +4,9 @@
 from electionguard.election import CiphertextElectionContext
 from electionguard.scheduler import Scheduler
 from electionguard.serializable import write_json_object
-from fastapi import APIRouter, Body
+from fastapi import APIRouter, Body, Depends
 
+from app.core.scheduler import get_scheduler
 from ..models import (
     convert_guardian,
     DecryptBallotSharesRequest,
@@ -17,7 +18,10 @@
 
 
 @router.post("/decrypt-shares", tags=[TALLY])
-def decrypt_ballot_shares(request: DecryptBallotSharesRequest = Body(...)) -> Any:
+def decrypt_ballot_shares(
+    request: DecryptBallotSharesRequest = Body(...),
+    scheduler: Scheduler = Depends(get_scheduler),
+) -> Any:
     """
     Decrypt this guardian's share of one or more ballots
     """
@@ -28,7 +32,6 @@ def decrypt_ballot_shares(request: DecryptBallotSharesRequest = Body(...)) -> An
     context = CiphertextElectionContext.from_json_object(request.context)
     guardian = convert_guardian(request.guardian)
 
-    scheduler = Scheduler()
     shares = [
         compute_decryption_share_for_ballot(guardian, ballot, context, scheduler)
         for ballot in ballots

app/api/v1/guardian/tally.py

@@ -5,10 +5,11 @@
     ElectionDescription,
     InternalElectionDescription,
 )
+from electionguard.scheduler import Scheduler
 from electionguard.serializable import write_json_object
-from fastapi import APIRouter, Body
-
+from fastapi import APIRouter, Body, Depends
 
+from app.core.scheduler import get_scheduler
 from ..models import (
     convert_guardian,
     convert_tally,
@@ -20,7 +21,10 @@
 
 
 @router.post("/decrypt-share", tags=[TALLY])
-def decrypt_share(request: DecryptTallyShareRequest = Body(...)) -> Any:
+def decrypt_share(
+    request: DecryptTallyShareRequest = Body(...),
+    scheduler: Scheduler = Depends(get_scheduler),
+) -> Any:
     """
     Decrypt a single guardian's share of a tally
     """
@@ -31,6 +35,6 @@ def decrypt_share(request: DecryptTallyShareRequest = Body(...)) -> Any:
     guardian = convert_guardian(request.guardian)
     tally = convert_tally(request.encrypted_tally, description, context)
 
-    share = compute_decryption_share(guardian, tally, context)
+    share = compute_decryption_share(guardian, tally, context, scheduler)
 
     return write_json_object(share)

app/api/v1/mediator/tally.py

@@ -7,15 +7,16 @@
     ElectionDescription,
     InternalElectionDescription,
 )
+from electionguard.scheduler import Scheduler
 from electionguard.serializable import read_json_object
 from electionguard.tally import (
     publish_ciphertext_tally,
     publish_plaintext_tally,
     CiphertextTally,
 )
-from fastapi import APIRouter, Body, HTTPException
-
+from fastapi import APIRouter, Body, Depends, HTTPException
 
+from app.core.scheduler import get_scheduler
 from ..models import (
     convert_tally,
     AppendTallyRequest,
@@ -29,27 +30,33 @@
 
 
 @router.post("", tags=[TALLY])
-def start_tally(request: StartTallyRequest = Body(...)) -> Any:
+def start_tally(
+    request: StartTallyRequest = Body(...),
+    scheduler: Scheduler = Depends(get_scheduler),
+) -> Any:
     """
     Start a new tally of a collection of ballots
     """
 
     ballots, description, context = _parse_tally_request(request)
     tally = CiphertextTally("election-results", description, context)
 
-    return _tally_ballots(tally, ballots)
+    return _tally_ballots(tally, ballots, scheduler)
 
 
 @router.post("/append", tags=[TALLY])
-def append_to_tally(request: AppendTallyRequest = Body(...)) -> Any:
+def append_to_tally(
+    request: AppendTallyRequest = Body(...),
+    scheduler: Scheduler = Depends(get_scheduler),
+) -> Any:
     """
     Append ballots into an existing tally
     """
 
     ballots, description, context = _parse_tally_request(request)
     tally = convert_tally(request.encrypted_tally, description, context)
 
-    return _tally_ballots(tally, ballots)
+    return _tally_ballots(tally, ballots, scheduler)
 
 
 @router.post("/decrypt", tags=[TALLY])
@@ -100,12 +107,14 @@ def _parse_tally_request(
 
 
 def _tally_ballots(
-    tally: CiphertextTally, ballots: List[CiphertextAcceptedBallot]
+    tally: CiphertextTally,
+    ballots: List[CiphertextAcceptedBallot],
+    scheduler: Scheduler,
 ) -> Any:
     """
     Append a series of ballots to a new or existing tally
     """
-    tally_succeeded = tally.batch_append(ballots)
+    tally_succeeded = tally.batch_append(ballots, scheduler)
 
     if tally_succeeded:
         published_tally = publish_ciphertext_tally(tally)

app/core/scheduler.py

@@ -0,0 +1,9 @@
+from functools import lru_cache
+from electionguard.scheduler import Scheduler
+
+__all__ = ["get_scheduler"]
+
+
+@lru_cache
+def get_scheduler() -> Scheduler:
+    return Scheduler()

app/main.py

@@ -5,6 +5,7 @@
 
 from app.api.v1.routes import get_routes
 from app.core.config import Settings
+from app.core.scheduler import get_scheduler
 
 logger = getLogger(__name__)
 
@@ -37,6 +38,14 @@ def get_app(settings: Optional[Settings] = None) -> FastAPI:
 
 app = get_app()
 
+
+@app.on_event("shutdown")
+def on_shutdown() -> None:
+    # Ensure a clean shutdown of the singleton Scheduler
+    scheduler = get_scheduler()
+    scheduler.close()
+
+
 if __name__ == "__main__":
     # IMPORTANT: This should only be used to debug the application.
     # For normal execution, run `make start`.

docker-compose.dev.yml

@@ -9,23 +9,23 @@ services:
       context: .
       target: dev
     volumes:
-      - "./app:/app"
+      - "./app:/app/app"
     ports:
       - 8000:8000
     environment:
       API_MODE: "mediator"
       PROJECT_NAME: "ElectionGuard Mediator API"
-      port: 8000
+      PORT: 8000
 
   guardian:
     build:
       context: .
       target: dev
     volumes:
-      - "./app:/app"
+      - "./app:/app/app"
     ports:
       - 8001:8001
     environment:
       API_MODE: "guardian"
       PROJECT_NAME: "ElectionGuard Guardian API"
-      port: 8001
+      PORT: 8001

docker-compose.yml

@@ -9,7 +9,7 @@ services:
     environment:
       API_MODE: "mediator"
       PROJECT_NAME: "ElectionGuard Mediator API"
-      port: 8000
+      PORT: 8000
 
   guardian:
     build:
@@ -20,4 +20,4 @@ services:
     environment:
       API_MODE: "guardian"
       PROJECT_NAME: "ElectionGuard Guardian API"
-      port: 8001
+      PORT: 8001

tests/integration/conftest.py

@@ -0,0 +1,15 @@
+from typing import Generator
+import pytest
+from app.core.scheduler import get_scheduler
+
+
+@pytest.yield_fixture(scope="session", autouse=True)
+def scheduler_lifespan() -> Generator[None, None, None]:
+    """
+    Ensure that the global scheduler singleton is
+    torn down when tests finish.  Otherwise, the test runner will hang
+    waiting for the scheduler to complete.
+    """
+    yield None
+    scheduler = get_scheduler()
+    scheduler.close()

tests/postman/docker-compose.yml

@@ -18,7 +18,7 @@ services:
     environment:
       API_MODE: "mediator"
       PROJECT_NAME: "ElectionGuard Mediator API"
-      port: 80
+      PORT: 80
 
   guardian:
     build:
@@ -29,7 +29,7 @@ services:
     environment:
       API_MODE: "guardian"
       PROJECT_NAME: "ElectionGuard Guardian API"
-      port: 80
+      PORT: 80
 
   test-runner:
     build: