Fix segments by repair run for memory storage

thelastpickle · Mar 26, 2024 · b378cac · b378cac
1 parent 4e7e5c9
commit b378cac
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 30 deletions.
diff --git a/src/server/src/main/java/io/cassandrareaper/storage/MemoryStorageFacade.java b/src/server/src/main/java/io/cassandrareaper/storage/MemoryStorageFacade.java
@@ -44,12 +44,11 @@
 import java.nio.file.Paths;
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
-import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.stream.Collectors;
 
 import org.eclipse.serializer.persistence.types.PersistenceFieldEvaluator;
 import org.eclipse.store.storage.embedded.types.EmbeddedStorage;
@@ -281,29 +280,13 @@ public RepairSegment addRepairSegment(RepairSegment segment) {
     LOG.info("Adding RepairSegment Token Ranges: {}",
         Arrays.toString(segment.getTokenRange().getTokenRanges().toArray()));
     final RepairSegment newSegment = this.memoryStorageRoot.addRepairSegment(segment);
-    // also add the segment by RunId
-    UUID repairSegmentRunId = segment.getRunId();
-    LinkedHashMap<UUID, RepairSegment> segmentsByRunId =
-        this.memoryStorageRoot.getRepairSegmentsByRunId().get(repairSegmentRunId);
-    if (segmentsByRunId == null) {
-      segmentsByRunId = new LinkedHashMap<>();
-      this.memoryStorageRoot.getRepairSegmentsByRunId().put(repairSegmentRunId, segmentsByRunId);
-    }
-    segmentsByRunId.put(segment.getId(), segment);
-    this.persist(this.memoryStorageRoot.getRepairSegments(), this.memoryStorageRoot.getRepairSegmentsByRunId());
+    this.persist(this.memoryStorageRoot.getRepairSegments());
     return newSegment;
   }
 
   public RepairSegment removeRepairSegment(UUID id) {
     RepairSegment segment = this.memoryStorageRoot.removeRepairSegment(id);
-    // also remove the segment from the byRunId map
-    UUID repairSegmentRunId = segment.getRunId();
-    LinkedHashMap<UUID, RepairSegment> segmentsByRunId =
-        this.memoryStorageRoot.getRepairSegmentsByRunId().get(repairSegmentRunId);
-    if (segmentsByRunId != null) {
-      segmentsByRunId.remove(segment.getId());
-    }
-    this.persist(this.memoryStorageRoot.getRepairSegments(), this.memoryStorageRoot.getRepairSegmentsByRunId());
+    this.persist(this.memoryStorageRoot.getRepairSegments());
     return segment;
   }
 
@@ -312,10 +295,8 @@ public RepairSegment getRepairSegmentById(UUID id) {
   }
 
   public Collection<RepairSegment> getRepairSegmentsByRunId(UUID runId) {
-    if (this.memoryStorageRoot.getRepairSegmentsByRunId().containsKey(runId)) {
-      return this.memoryStorageRoot.getRepairSegmentsByRunId().get(runId).values();
-    }
-    return Collections.EMPTY_LIST;
+    return this.memoryStorageRoot.getRepairSegments().values().stream()
+        .filter(segment -> segment.getRunId().equals(runId)).collect(Collectors.toSet());
   }
 
   // RepairSubscription operations

diff --git a/src/server/src/main/java/io/cassandrareaper/storage/memory/MemoryStorageRoot.java b/src/server/src/main/java/io/cassandrareaper/storage/memory/MemoryStorageRoot.java
@@ -23,7 +23,6 @@
 import io.cassandrareaper.core.RepairSegment;
 import io.cassandrareaper.core.RepairUnit;
 
-import java.util.LinkedHashMap;
 import java.util.Map;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentMap;
@@ -32,7 +31,6 @@
 
 
 public final class MemoryStorageRoot {
-  private final ConcurrentMap<UUID, LinkedHashMap<UUID, RepairSegment>> repairSegmentsByRunId = Maps.newConcurrentMap();
   private final ConcurrentMap<UUID, RepairSegment> repairSegments = Maps.newConcurrentMap();
   private final ConcurrentMap<UUID, RepairUnit> repairUnits = Maps.newConcurrentMap();
   private final ConcurrentMap<RepairUnit.Builder, RepairUnit> repairUnitsByKey = Maps.newConcurrentMap();
@@ -76,10 +74,6 @@ public RepairSchedule removeRepairSchedule(UUID id) {
   }
 
   // RepairSegment operations
-  public Map<UUID, LinkedHashMap<UUID, RepairSegment>> getRepairSegmentsByRunId() {
-    return this.repairSegmentsByRunId;
-  }
-
   public Map<UUID, RepairSegment> getRepairSegments() {
     return this.repairSegments;
   }