[dpc-4710] Shuts down dpc-queue after AggregationEngine (#2708)

## 🎫 Ticket

https://jira.cms.gov/browse/DPC-4710
Bugfix for #2706

## 🛠 Changes

Updated the order that `DPCQueueHibernateModule` and
`AggregationAppModule` are added to the `DPCAggregationService`.

## ℹ️ Context

DropWizard shuts these modules down in the opposite order they are
added, and we don't want the queue to shutdown until after the
aggregation engine has finished processing or it won't be able to update
its last batch.

This wasn't caught in initial testing for two reasons:

1. Our tests used the `MemoryBatchQueue` which doesn't connect to an
external DB.
2. Our code was ignoring the errors when it tried to pause a batch
instead of logging them.

## 🧪 Validation

1. Go into your local queue DB and pick out your most recently completed
batch.
2. Update the batch's `status` to 0 (_queued_), its `aggregator_id` to
`null`, and its patient to `9S99EU8XY92`. (this patient tells
MockBfdClient to hang forever, simulating a long running patient)
3. Start up dpc-aggregation and it will pick up the batch and set its
status to 1 (_running_).
4. Stop dpc-aggregation.  

If you do this on the current main branch, the status of the batch will
still be 1 (_running_). If you recompile from this branch and go through
the same process the batch will be reset to 0 (_queued_), which is what
we want.

Author: MEspositoE14s

dpc-aggregation/src/main/java/gov/cms/dpc/aggregation/DPCAggregationService.java

@@ -42,9 +42,12 @@ public void initialize(Bootstrap<DPCAggregationConfiguration> bootstrap) {
                 new SubstitutingSourceProvider(bootstrap.getConfigurationSourceProvider(), substitutor);
         bootstrap.setConfigurationSourceProvider(provider);
 
+        // AggregationAppModule needs to be added after DPCQueueHibernateModule, because DropWizard shuts these down
+        // in reverse order, and we don't want the queue DB to be disconnected until after the aggregation engine stops
+        // running.
         GuiceBundle guiceBundle = GuiceBundle.builder()
-                .modules(new AggregationAppModule(),
-                        new DPCQueueHibernateModule<>(queueHibernateBundle),
+                .modules(new DPCQueueHibernateModule<>(queueHibernateBundle),
+                        new AggregationAppModule(),
                         new DPCHibernateModule<>(hibernateBundle),
                         new JobQueueModule<DPCAggregationConfiguration>(),
                         new BlueButtonClientModule<DPCAggregationConfiguration>())

dpc-aggregation/src/main/java/gov/cms/dpc/aggregation/engine/AggregationEngine.java

@@ -96,7 +96,10 @@ public void stop() {
 
         // If a batch is currently running, mark it paused.
         Optional<JobQueueBatch> optionalBatch = this.currentBatch.get();
-        optionalBatch.ifPresent(jobQueueBatch -> this.queue.pauseBatch(jobQueueBatch, aggregatorID));
+        optionalBatch.ifPresent(jobQueueBatch -> {
+            logger.info("Pausing batch: {}", jobQueueBatch.getBatchID());
+            this.queue.pauseBatch(jobQueueBatch, aggregatorID);
+        });
     }
 
     public boolean isRunning() {

dpc-aggregation/src/test/java/gov/cms/dpc/aggregation/AggregationServiceTest.java

@@ -1,19 +1,32 @@
 package gov.cms.dpc.aggregation;
 
+import com.codahale.metrics.MetricRegistry;
 import com.codahale.metrics.health.HealthCheckRegistry;
+import gov.cms.dpc.bluebutton.client.MockBlueButtonClient;
+import gov.cms.dpc.common.hibernate.queue.DPCQueueManagedSessionFactory;
+import gov.cms.dpc.common.utils.NPIUtil;
+import gov.cms.dpc.fhir.DPCResourceType;
+import gov.cms.dpc.queue.DistributedBatchQueue;
+import gov.cms.dpc.queue.JobStatus;
+import gov.cms.dpc.queue.models.JobQueueBatch;
 import gov.cms.dpc.testing.BufferedLoggerHandler;
 import gov.cms.dpc.testing.IntegrationTest;
 import io.dropwizard.testing.ConfigOverride;
 import io.dropwizard.testing.DropwizardTestSupport;
-import org.junit.jupiter.api.AfterAll;
-import org.junit.jupiter.api.BeforeAll;
+import org.hibernate.SessionFactory;
+import org.hibernate.cfg.Configuration;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.SortedSet;
+import java.util.UUID;
 
-import static org.junit.jupiter.api.Assertions.assertAll;
-import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.awaitility.Awaitility.await;
+import static org.junit.jupiter.api.Assertions.*;
 
 @IntegrationTest
 @ExtendWith(BufferedLoggerHandler.class)
@@ -25,13 +38,19 @@ public class AggregationServiceTest {
                     ConfigOverride.config("server.applicationConnectors[0].port", "7777"),
                     ConfigOverride.config("server.adminConnectors[0].port", "8888"));
 
-    @BeforeAll
-    static void start() throws Exception{
+	private DistributedBatchQueue queue;
+
+    @BeforeEach
+    void start() throws Exception{
+        final Configuration conf = new Configuration();
+		SessionFactory sessionFactory = conf.configure().buildSessionFactory();
+        queue = new DistributedBatchQueue(new DPCQueueManagedSessionFactory(sessionFactory), 100, new MetricRegistry());
+
         APPLICATION.before();
     }
 
-    @AfterAll
-    static void stop() {
+    @AfterEach
+    void stop() {
         APPLICATION.after();
     }
 
@@ -48,4 +67,31 @@ void testHealthChecks() {
         // Everything should be true
         checks.runHealthChecks().forEach((key, value) -> assertTrue(value.isHealthy(), String.format("Healthcheck: %s is not ok.", key)));
     }
+
+    @Test
+    void testStoppingEngineMidBatch() throws InterruptedException {
+        // Create a batch that will hang forever
+        final var orgID = UUID.randomUUID();
+        final List<String> mbis = List.of(MockBlueButtonClient.TEST_PATIENT_TIME_OUT);
+
+        final var jobID = queue.createJob(
+            orgID,
+            NPIUtil.generateNPI(),
+            NPIUtil.generateNPI(),
+            mbis,
+            Collections.singletonList(DPCResourceType.Patient),
+            null,
+            MockBlueButtonClient.BFD_TRANSACTION_TIME,
+            null, null, true, false);
+
+        // Wait for dpc-aggregation to pick up the batch, set status to RUNNING and hang on getting the patient from BFD
+        await().until(() -> {
+            JobQueueBatch batch = queue.getJobBatches(jobID).get(0);
+            return batch.getStatus() == JobStatus.RUNNING;
+        });
+
+        // Stop aggregation and make sure it pauses the batch
+        APPLICATION.after();
+        assertEquals(JobStatus.QUEUED, queue.getJobBatches(jobID).get(0).getStatus());
+    }
 }

dpc-aggregation/src/test/resources/hibernate.cfg.xml

@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!DOCTYPE hibernate-configuration PUBLIC
+        "-//Hibernate/Hibernate Configuration DTD 3.0//EN"
+        "http://www.hibernate.org/dtd/hibernate-configuration-3.0.dtd">
+
+<hibernate-configuration>
+    <session-factory>
+
+        <!-- Connection settings -->
+        <property name="hibernate.connection.driver_class">org.postgresql.Driver</property>
+        <property name="hibernate.connection.url">jdbc:postgresql://localhost:5432/dpc_queue</property>
+        <property name="hibernate.connection.username">postgres</property>
+        <property name="hibernate.connection.password">dpc-safe</property>
+
+        <!-- SQL dialect -->
+        <property name="hibernate.dialect">org.hibernate.dialect.PostgreSQLDialect</property>
+
+        <!-- Print executed SQL to stdout -->
+        <property name="show_sql">false</property>
+
+        <!-- Update database on startup -->
+        <property name="hibernate.hbm2ddl.auto">validate</property>
+        <property name="hibernate.connection.autocommit">true</property>
+
+        <!-- Annotated entity classes -->
+        <mapping class="gov.cms.dpc.queue.models.JobQueueBatch"/>
+        <mapping class="gov.cms.dpc.queue.models.JobQueueBatchFile"/>
+
+    </session-factory>
+</hibernate-configuration>

dpc-queue/src/main/java/gov/cms/dpc/queue/DistributedBatchQueue.java

@@ -239,6 +239,8 @@ public void pauseBatch(JobQueueBatch job, UUID aggregatorID) {
             } finally {
                 tx.commit();
             }
+        } catch(Exception e) {
+            logger.error("Error pausing batch: {} {}", job.getBatchID(), e.getMessage());
         }
     }