From b3ce95ac3b8ea65287628486f4eaff6a942b0b75 Mon Sep 17 00:00:00 2001
From: DimaDemchenko <demchenkodmytro73@gmail.com>
Date: Thu, 19 Sep 2024 11:10:07 +0300
Subject: [PATCH] refactor: Improve clear logic and added getAvailableSpace
 func

---
 .../src/segment-storage/index.ts              |   8 ++
 .../segment-storage/segment-memory-storage.ts | 111 +++++++++++-------
 2 files changed, 74 insertions(+), 45 deletions(-)

diff --git a/packages/p2p-media-loader-core/src/segment-storage/index.ts b/packages/p2p-media-loader-core/src/segment-storage/index.ts
index bf82fe3b..8aa7c569 100644
--- a/packages/p2p-media-loader-core/src/segment-storage/index.ts
+++ b/packages/p2p-media-loader-core/src/segment-storage/index.ts
@@ -74,6 +74,14 @@ export interface SegmentStorage {
     swarmId: string,
   ): Promise<ArrayBuffer | undefined>;
 
+  /**
+   * Returns available space in the storage
+   */
+  getAvailableSpace(): {
+    limit: number;
+    used: number;
+  };
+
   /**
    * Returns true if segment is in storage
    * @param streamId - Stream identifier
diff --git a/packages/p2p-media-loader-core/src/segment-storage/segment-memory-storage.ts b/packages/p2p-media-loader-core/src/segment-storage/segment-memory-storage.ts
index 9fd795f7..c1df99d5 100644
--- a/packages/p2p-media-loader-core/src/segment-storage/segment-memory-storage.ts
+++ b/packages/p2p-media-loader-core/src/segment-storage/segment-memory-storage.ts
@@ -21,12 +21,11 @@ type LastRequestedSegmentInfo = {
   segmentId: number;
   startTime: number;
   endTime: number;
+  swarmId: string;
+  streamType: StreamType;
+  isLiveStream: boolean;
 };
 
-type SegmentCategories = {
-  obsolete: string[];
-  aheadHttpWindow: string[];
-};
 const getStorageItemId = (streamId: string, segmentId: number) =>
   `${streamId}|${segmentId}`;
 
@@ -81,15 +80,18 @@ export class SegmentMemoryStorage implements SegmentStorage {
     segmentId: number,
     startTime: number,
     endTime: number,
-    _swarmId: string,
-    _streamType: StreamType,
-    _isLiveStream: boolean,
+    swarmId: string,
+    streamType: StreamType,
+    isLiveStream: boolean,
   ): void {
     this.lastRequestedSegment = {
       streamId,
       segmentId,
       startTime,
       endTime,
+      swarmId,
+      streamType,
+      isLiveStream,
     };
   }
 
@@ -136,6 +138,30 @@ export class SegmentMemoryStorage implements SegmentStorage {
     return dataItem.data;
   }
 
+  getAvailableSpace() {
+    if (!this.lastRequestedSegment) {
+      return {
+        limit: this.segmentsMemoryStorageLimit,
+        used: this.currentMemoryStorageSize,
+      };
+    }
+    const { streamId, isLiveStream } = this.lastRequestedSegment;
+    const segmentsToRemove = this.findSegmentsToRemove(isLiveStream, streamId);
+
+    const potentialFreeSpace = segmentsToRemove.reduce((total, segmentId) => {
+      const segment = this.cache.get(segmentId);
+      return segment ? total + segment.data.byteLength : total;
+    }, 0);
+
+    const usedMemoryInMB =
+      this.currentMemoryStorageSize - potentialFreeSpace / BYTES_PER_MB;
+
+    return {
+      limit: this.segmentsMemoryStorageLimit,
+      used: usedMemoryInMB,
+    };
+  }
+
   hasSegment(streamId: string, externalId: number, _swarmId: string) {
     const segmentStorageId = getStorageItemId(streamId, externalId);
     const segment = this.cache.get(segmentStorageId);
@@ -241,54 +267,49 @@ export class SegmentMemoryStorage implements SegmentStorage {
       return [];
     }
 
-    const segmentsToRemove: SegmentCategories = {
-      obsolete: [],
-      aheadHttpWindow: [],
-    };
-
-    const currentPlayback = this.currentPlayback.position;
-    const sortedCache = Array.from(this.cache.values()).sort(
-      (a, b) => a.endTime - b.endTime,
-    );
+    const obsoleteSegments: string[] = [];
+    const playbackPosition = this.currentPlayback.position;
 
-    for (const segmentData of sortedCache) {
-      const { streamId, segmentId, endTime, streamType } = segmentData;
+    for (const segmentData of this.cache.values()) {
+      const { streamId, segmentId } = segmentData;
       const storageId = getStorageItemId(streamId, segmentId);
 
-      if (streamId !== currentStreamId) {
-        segmentsToRemove.obsolete.push(storageId);
-      }
-
-      const highDemandTimeWindow = this.getStreamTimeWindow(
-        streamType,
-        "highDemandTimeWindow",
-      );
-      const httpDownloadTimeWindow = this.getStreamTimeWindow(
-        streamType,
-        "httpDownloadTimeWindow",
+      const shouldRemove = this.shouldRemoveSegment(
+        segmentData,
+        isLiveStream,
+        currentStreamId,
+        playbackPosition,
       );
 
-      if (isLiveStream && currentPlayback > highDemandTimeWindow + endTime) {
-        segmentsToRemove.obsolete.push(storageId);
-        continue;
-      }
+      if (shouldRemove) obsoleteSegments.push(storageId);
+    }
 
-      if (currentPlayback > endTime) {
-        segmentsToRemove.obsolete.push(storageId);
-      }
+    return obsoleteSegments;
+  }
 
-      if (segmentsToRemove.obsolete.length > 0) {
-        break;
-      }
+  private shouldRemoveSegment(
+    segmentData: SegmentDataItem,
+    isLiveStream: boolean,
+    currentStreamId: string,
+    currentPlaybackPosition: number,
+  ): boolean {
+    const { streamId, endTime, streamType } = segmentData;
+    const highDemandTimeWindow = this.getStreamTimeWindow(
+      streamType,
+      "highDemandTimeWindow",
+    );
 
-      if (endTime > currentPlayback + httpDownloadTimeWindow) {
-        segmentsToRemove.aheadHttpWindow.push(storageId);
-      }
+    if (currentPlaybackPosition <= endTime) return false;
+    if (streamId !== currentStreamId) return true;
+    if (
+      isLiveStream &&
+      currentPlaybackPosition > highDemandTimeWindow + endTime
+    ) {
+      return true;
     }
+    if (!isLiveStream) return true;
 
-    return segmentsToRemove.obsolete.length > 0
-      ? segmentsToRemove.obsolete
-      : segmentsToRemove.aheadHttpWindow;
+    return false;
   }
 
   private setMemoryStorageLimit() {