8
8
9
9
package org .opensearch .index .remote ;
10
10
11
+ import org .apache .logging .log4j .Logger ;
12
+ import org .apache .logging .log4j .message .ParameterizedMessage ;
13
+ import org .opensearch .common .CheckedFunction ;
14
+ import org .opensearch .common .logging .Loggers ;
11
15
import org .opensearch .core .common .io .stream .StreamInput ;
12
16
import org .opensearch .core .common .io .stream .StreamOutput ;
13
17
import org .opensearch .core .common .io .stream .Writeable ;
17
21
import org .opensearch .core .index .shard .ShardId ;
18
22
19
23
import java .io .IOException ;
20
- import java .util .HashMap ;
24
+ import java .util .Collection ;
25
+ import java .util .Collections ;
21
26
import java .util .HashSet ;
22
27
import java .util .Map ;
23
28
import java .util .Set ;
24
29
import java .util .concurrent .atomic .AtomicLong ;
25
30
import java .util .concurrent .atomic .AtomicReference ;
26
31
import java .util .stream .Collectors ;
27
32
33
+ import static org .opensearch .index .shard .RemoteStoreRefreshListener .EXCLUDE_FILES ;
34
+
28
35
/**
29
36
* Keeps track of remote refresh which happens in {@link org.opensearch.index.shard.RemoteStoreRefreshListener}. This consist of multiple critical metrics.
30
37
*
31
38
* @opensearch.internal
32
39
*/
33
40
public class RemoteRefreshSegmentTracker {
34
41
42
+ private final Logger logger ;
43
+
35
44
/**
36
45
* ShardId for which this instance tracks the remote segment upload metadata.
37
46
*/
@@ -123,14 +132,14 @@ public class RemoteRefreshSegmentTracker {
123
132
private final Map <String , AtomicLong > rejectionCountMap = ConcurrentCollections .newConcurrentMap ();
124
133
125
134
/**
126
- * Map of name to size of the segment files created as part of the most recent refresh.
135
+ * Keeps track of segment files and their size in bytes which are part of the most recent refresh.
127
136
*/
128
- private volatile Map <String , Long > latestLocalFileNameLengthMap ;
137
+ private final Map <String , Long > latestLocalFileNameLengthMap = ConcurrentCollections . newConcurrentMap () ;
129
138
130
139
/**
131
140
* Set of names of segment files that were uploaded as part of the most recent remote refresh.
132
141
*/
133
- private final Set <String > latestUploadedFiles = new HashSet <> ();
142
+ private final Set <String > latestUploadedFiles = ConcurrentCollections . newConcurrentSet ();
134
143
135
144
/**
136
145
* Keeps the bytes lag computed so that we do not compute it for every request.
@@ -175,6 +184,7 @@ public RemoteRefreshSegmentTracker(
175
184
int uploadBytesPerSecMovingAverageWindowSize ,
176
185
int uploadTimeMsMovingAverageWindowSize
177
186
) {
187
+ logger = Loggers .getLogger (getClass (), shardId );
178
188
this .shardId = shardId ;
179
189
// Both the local refresh time and remote refresh time are set with current time to give consistent view of time lag when it arises.
180
190
long currentClockTimeMs = System .currentTimeMillis ();
@@ -186,8 +196,6 @@ public RemoteRefreshSegmentTracker(
186
196
uploadBytesMovingAverageReference = new AtomicReference <>(new MovingAverage (uploadBytesMovingAverageWindowSize ));
187
197
uploadBytesPerSecMovingAverageReference = new AtomicReference <>(new MovingAverage (uploadBytesPerSecMovingAverageWindowSize ));
188
198
uploadTimeMsMovingAverageReference = new AtomicReference <>(new MovingAverage (uploadTimeMsMovingAverageWindowSize ));
189
-
190
- latestLocalFileNameLengthMap = new HashMap <>();
191
199
}
192
200
193
201
ShardId getShardId () {
@@ -361,12 +369,36 @@ long getRejectionCount(String rejectionReason) {
361
369
return rejectionCountMap .get (rejectionReason ).get ();
362
370
}
363
371
364
- Map <String , Long > getLatestLocalFileNameLengthMap () {
365
- return latestLocalFileNameLengthMap ;
372
+ public Map <String , Long > getLatestLocalFileNameLengthMap () {
373
+ return Collections . unmodifiableMap ( latestLocalFileNameLengthMap ) ;
366
374
}
367
375
368
- public void setLatestLocalFileNameLengthMap (Map <String , Long > latestLocalFileNameLengthMap ) {
369
- this .latestLocalFileNameLengthMap = latestLocalFileNameLengthMap ;
376
+ /**
377
+ * Updates the latestLocalFileNameLengthMap by adding file name and it's size to the map. The method is given a function as an argument which is used for determining the file size (length in bytes). This method is also provided the collection of segment files which are the latest refresh local segment files. This method also removes the stale segment files from the map that are not part of the input segment files.
378
+ *
379
+ * @param segmentFiles list of local refreshed segment files
380
+ * @param fileSizeFunction function is used to determine the file size in bytes
381
+ */
382
+ public void updateLatestLocalFileNameLengthMap (
383
+ Collection <String > segmentFiles ,
384
+ CheckedFunction <String , Long , IOException > fileSizeFunction
385
+ ) {
386
+ // Update the map
387
+ segmentFiles .stream ()
388
+ .filter (file -> EXCLUDE_FILES .contains (file ) == false )
389
+ .filter (file -> latestLocalFileNameLengthMap .containsKey (file ) == false || latestLocalFileNameLengthMap .get (file ) == 0 )
390
+ .forEach (file -> {
391
+ long fileSize = 0 ;
392
+ try {
393
+ fileSize = fileSizeFunction .apply (file );
394
+ } catch (IOException e ) {
395
+ logger .warn (new ParameterizedMessage ("Exception while reading the fileLength of file={}" , file ), e );
396
+ }
397
+ latestLocalFileNameLengthMap .put (file , fileSize );
398
+ });
399
+ Set <String > fileSet = new HashSet <>(segmentFiles );
400
+ // Remove keys from the fileSizeMap that do not exist in the latest segment files
401
+ latestLocalFileNameLengthMap .entrySet ().removeIf (entry -> fileSet .contains (entry .getKey ()) == false );
370
402
computeBytesLag ();
371
403
}
372
404
@@ -382,7 +414,7 @@ public void setLatestUploadedFiles(Set<String> files) {
382
414
}
383
415
384
416
private void computeBytesLag () {
385
- if (latestLocalFileNameLengthMap == null || latestLocalFileNameLengthMap .isEmpty ()) {
417
+ if (latestLocalFileNameLengthMap .isEmpty ()) {
386
418
return ;
387
419
}
388
420
Set <String > filesNotYetUploaded = latestLocalFileNameLengthMap .keySet ()
0 commit comments