Mimic storage notifications using the HTTP API (#1339)

* Add stubbed polling of job details from the mgmt-api. This will not work without the actual client implementation Implement using apiClient the triggerRepair, getJobDetails, scheduler as well as add a simple test to ensure the state is managed correctly * Merge test files after the rebase * Add a test to verify the behavior of the notifications polling * Address comments
thelastpickle · Aug 28, 2023 · e5bbf2f · e5bbf2f
1 parent d8497ac
commit e5bbf2f
Show file tree

Hide file tree

Showing 6 changed files with 268 additions and 14 deletions.
diff --git a/src/server/pom.xml b/src/server/pom.xml
@@ -274,7 +274,7 @@
         <dependency>
             <groupId>io.k8ssandra</groupId>
             <artifactId>datastax-mgmtapi-client-openapi</artifactId>
-            <version>0.1.0-c22a2fc</version>
+            <version>0.1.0-4d2a772</version>
         </dependency>
         <!--test scope -->
 

diff --git a/src/server/src/main/java/io/cassandrareaper/ReaperApplication.java b/src/server/src/main/java/io/cassandrareaper/ReaperApplication.java
@@ -179,7 +179,7 @@ public void run(ReaperApplicationConfiguration config, Environment environment)
     Cryptograph cryptograph = context.config == null || context.config.getCryptograph() == null
         ? new NoopCrypotograph() : context.config.getCryptograph().create();
 
-    initializeManagement(context, cryptograph);
+    initializeManagement(context, environment, cryptograph);
 
     context.repairManager = RepairManager.create(
         context,
@@ -299,14 +299,17 @@ public void run(ReaperApplicationConfiguration config, Environment environment)
   }
 
 
-  private void initializeManagement(AppContext context, Cryptograph cryptograph) {
+  private void initializeManagement(AppContext context, Environment environment, Cryptograph cryptograph) {
     if (context.managementConnectionFactory == null) {
       LOG.info("no management connection factory given in context, creating default");
       if (context.config.getHttpManagement() == null || !context.config.getHttpManagement().isEnabled()) {
         LOG.info("HTTP management connection config not set, or set disabled. Creating JMX connection factory instead");
         context.managementConnectionFactory = new JmxManagementConnectionFactory(context, cryptograph);
       } else {
-        context.managementConnectionFactory = new HttpManagementConnectionFactory(context);
+        ScheduledExecutorService jobStatusPollerExecutor = environment.lifecycle()
+            .scheduledExecutorService("JobStatusPoller")
+            .threads(2).build();
+        context.managementConnectionFactory = new HttpManagementConnectionFactory(context, jobStatusPollerExecutor);
       }
     }
   }

diff --git a/...server/src/main/java/io/cassandrareaper/management/http/HttpCassandraManagementProxy.java b/...server/src/main/java/io/cassandrareaper/management/http/HttpCassandraManagementProxy.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2020 The Last Pickle Ltd
+ * Copyright 2023-2023 DataStax, Inc.
  *
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -22,6 +22,7 @@
 import io.cassandrareaper.core.Table;
 import io.cassandrareaper.management.ICassandraManagementProxy;
 import io.cassandrareaper.management.RepairStatusHandler;
+import io.cassandrareaper.management.http.models.JobStatusTracker;
 import io.cassandrareaper.service.RingRange;
 
 import java.io.IOException;
@@ -35,42 +36,66 @@
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
 import javax.management.JMException;
 import javax.management.openmbean.CompositeData;
 import javax.validation.constraints.NotNull;
 
 import com.codahale.metrics.MetricRegistry;
 import com.datastax.mgmtapi.client.api.DefaultApi;
-import com.datastax.mgmtapi.client.invoker.ApiClient;
 import com.datastax.mgmtapi.client.invoker.ApiException;
 import com.datastax.mgmtapi.client.model.EndpointStates;
+import com.datastax.mgmtapi.client.model.Job;
+import com.datastax.mgmtapi.client.model.RepairRequest;
 import com.datastax.mgmtapi.client.model.SnapshotDetails;
+import com.datastax.mgmtapi.client.model.StatusChange;
 import com.datastax.mgmtapi.client.model.TakeSnapshotRequest;
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.Maps;
 import org.apache.cassandra.repair.RepairParallelism;
+import org.apache.cassandra.utils.progress.ProgressEventType;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public class HttpCassandraManagementProxy implements ICassandraManagementProxy {
+
+  public static final int DEFAULT_POLL_INTERVAL_IN_MILLISECONDS = 5000;
   private static final Logger LOG = LoggerFactory.getLogger(HttpCassandraManagementProxy.class);
   final String host;
   final MetricRegistry metricRegistry;
   final String rootPath;
   final InetSocketAddress endpoint;
   final DefaultApi apiClient;
 
+  final ConcurrentMap<Integer, RepairStatusHandler> repairStatusHandlers = Maps.newConcurrentMap();
+  final ConcurrentMap<String, JobStatusTracker> jobTracker = Maps.newConcurrentMap();
+  final ConcurrentMap<Integer, ExecutorService> repairStatusExecutors = Maps.newConcurrentMap();
+
+
+  private ScheduledExecutorService statusTracker;
+
   public HttpCassandraManagementProxy(MetricRegistry metricRegistry,
                                       String rootPath,
-                                      InetSocketAddress endpoint
+                                      InetSocketAddress endpoint,
+                                      ScheduledExecutorService executor,
+                                      DefaultApi apiClient
   ) {
     this.host = endpoint.getHostString();
     this.metricRegistry = metricRegistry;
     this.rootPath = rootPath;
     this.endpoint = endpoint;
-    this.apiClient = new DefaultApi(
-        new ApiClient().setBasePath("http://" + endpoint.getHostName() + ":" + endpoint.getPort() + rootPath));
+    this.apiClient = apiClient;
+    this.statusTracker = executor;
+
+    // TODO Perhaps the poll interval should be configurable through context.config ?
+    this.scheduleJobPoller(DEFAULT_POLL_INTERVAL_IN_MILLISECONDS);
   }
 
   @Override
@@ -192,13 +217,31 @@ public int triggerRepair(
       List<RingRange> associatedTokens,
       int repairThreadCount)
       throws ReaperException {
-    return 1; //TODO: implement me
 
+    String jobId;
+    try {
+      jobId = apiClient.repair1(new RepairRequest());
+    } catch (ApiException e) {
+      throw new ReaperException(e);
+    }
+
+    int repairNo = Integer.parseInt(jobId.substring(7));
+
+    repairStatusExecutors.putIfAbsent(repairNo, Executors.newSingleThreadExecutor());
+    repairStatusHandlers.putIfAbsent(repairNo, repairStatusHandler);
+    jobTracker.put(jobId, new JobStatusTracker());
+    return repairNo;
   }
 
   @Override
   public void removeRepairStatusHandler(int repairNo) {
-    // TODO: implement me.
+    repairStatusHandlers.remove(repairNo);
+    ExecutorService repairStatusExecutor = repairStatusExecutors.remove(repairNo);
+    if (null != repairStatusExecutor) {
+      repairStatusExecutor.shutdown();
+    }
+    String jobId = String.format("repair-%d", repairNo);
+    jobTracker.remove(jobId);
   }
 
   @Override
@@ -348,4 +391,54 @@ public String getUntranslatedHost() {
     //TODO: implement me
     return "";
   }
+
+  private Job getJobStatus(String id) {
+    // Poll with HTTP client the job's status
+    try {
+      return apiClient.getJobStatus(id);
+    } catch (ApiException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  @VisibleForTesting
+  private void scheduleJobPoller(int pollInterval) {
+    statusTracker.scheduleWithFixedDelay(
+        notificationsTracker(),
+        pollInterval * 2,
+        pollInterval,
+        TimeUnit.MILLISECONDS);
+  }
+
+  @VisibleForTesting
+  Runnable notificationsTracker() {
+    return () -> {
+      if (jobTracker.size() > 0) {
+        for (Map.Entry<String, JobStatusTracker> entry : jobTracker.entrySet()) {
+          Job job = getJobStatus(entry.getKey());
+          int availableNotifications = job.getStatusChanges().size();
+          int currentNotificationCount = entry.getValue().latestNotificationCount.get();
+
+          if (currentNotificationCount < availableNotifications) {
+            // We need to process the new ones
+            for (int i = currentNotificationCount; i < availableNotifications; i++) {
+              StatusChange statusChange = job.getStatusChanges().get(i);
+              // remove "repair-" prefix
+              int repairNo = Integer.parseInt(job.getId().substring(7));
+              ProgressEventType progressType = ProgressEventType.valueOf(statusChange.getStatus());
+              repairStatusExecutors.get(repairNo).submit(() -> {
+                repairStatusHandlers
+                    .get(repairNo)
+                    .handle(repairNo, Optional.empty(), Optional.of(progressType),
+                        statusChange.getMessage(), this);
+              });
+
+              // Update the count as we process them
+              entry.getValue().latestNotificationCount.incrementAndGet();
+            }
+          }
+        }
+      }
+    };
+  }
 }
diff --git a/...ver/src/main/java/io/cassandrareaper/management/http/HttpManagementConnectionFactory.java b/...ver/src/main/java/io/cassandrareaper/management/http/HttpManagementConnectionFactory.java
@@ -31,10 +31,14 @@
 import java.util.List;
 import java.util.Set;
 import java.util.concurrent.ConcurrentMap;
+import java.util.concurrent.ScheduledExecutorService;
 import javax.ws.rs.core.Response;
 
 import com.codahale.metrics.Gauge;
+import com.codahale.metrics.InstrumentedScheduledExecutorService;
 import com.codahale.metrics.MetricRegistry;
+import com.datastax.mgmtapi.client.api.DefaultApi;
+import com.datastax.mgmtapi.client.invoker.ApiClient;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.Maps;
 import com.google.common.collect.Sets;
@@ -47,14 +51,17 @@ public class HttpManagementConnectionFactory implements IManagementConnectionFac
   private final MetricRegistry metricRegistry;
   private final HostConnectionCounters hostConnectionCounters;
 
+  private ScheduledExecutorService jobStatusPollerExecutor;
+
   private final Set<String> accessibleDatacenters = Sets.newHashSet();
 
   // Constructor for HttpManagementConnectionFactory
-  public HttpManagementConnectionFactory(AppContext context) {
+  public HttpManagementConnectionFactory(AppContext context, ScheduledExecutorService jobStatusPollerExecutor) {
     this.metricRegistry
         = context.metricRegistry == null ? new MetricRegistry() : context.metricRegistry;
     hostConnectionCounters = new HostConnectionCounters(metricRegistry);
     registerConnectionsGauge();
+    this.jobStatusPollerExecutor = jobStatusPollerExecutor;
   }
 
   public ICassandraManagementProxy connectAny(Collection<Node> nodes) throws ReaperException {
@@ -118,10 +125,17 @@ private ICassandraManagementProxy connectImpl(Node node)
     if (pidResponse.getStatus() != 200) {
       throw new ReaperException("Could not get PID for node " + node.getHostname());
     }
+    DefaultApi apiClient = new DefaultApi(
+        new ApiClient().setBasePath("https://" + node.getHostname() + ":" + managementPort + rootPath));
+
+    InstrumentedScheduledExecutorService statusTracker = new InstrumentedScheduledExecutorService(
+        jobStatusPollerExecutor, metricRegistry);
     return new HttpCassandraManagementProxy(
         metricRegistry,
         rootPath,
-        new InetSocketAddress(node.getHostname(), managementPort)
+        new InetSocketAddress(node.getHostname(), managementPort),
+        statusTracker,
+        apiClient
     );
   }
 

diff --git a/src/server/src/main/java/io/cassandrareaper/management/http/models/JobStatusTracker.java b/src/server/src/main/java/io/cassandrareaper/management/http/models/JobStatusTracker.java
@@ -0,0 +1,23 @@
+/*
+ * Copyright 2023-2023 DataStax, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.cassandrareaper.management.http.models;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+public class JobStatusTracker {
+  public AtomicInteger latestNotificationCount = new AtomicInteger(0);
+}