Skip to content

Commit 0710a9d

Browse files
authored
Fix potential file leak in ES816BinaryQuantizedVectorsWriter (#120014) (#120091)
We are creating tmp files that might not get closed if an exception happens just after it. This commit makes sure all errors are handle properly and files are getting closed and deleted. # Conflicts: # muted-tests.yml
1 parent 7141ff1 commit 0710a9d

File tree

2 files changed

+55
-32
lines changed

2 files changed

+55
-32
lines changed

docs/changelog/120014.yaml

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
pr: 120014
2+
summary: Fix potential file leak in ES816BinaryQuantizedVectorsWriter
3+
area: Search
4+
type: bug
5+
issues:
6+
- 119981

server/src/main/java/org/elasticsearch/index/codec/vectors/ES816BinaryQuantizedVectorsWriter.java

Lines changed: 49 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -433,46 +433,66 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
433433
float[] centroid,
434434
float cDotC
435435
) throws IOException {
436-
long vectorDataOffset = binarizedVectorData.alignFilePointer(Float.BYTES);
437-
final IndexOutput tempQuantizedVectorData = segmentWriteState.directory.createTempOutput(
438-
binarizedVectorData.getName(),
439-
"temp",
440-
segmentWriteState.context
441-
);
442-
final IndexOutput tempScoreQuantizedVectorData = segmentWriteState.directory.createTempOutput(
443-
binarizedVectorData.getName(),
444-
"score_temp",
445-
segmentWriteState.context
446-
);
447-
IndexInput binarizedDataInput = null;
448-
IndexInput binarizedScoreDataInput = null;
449-
boolean success = false;
450-
int descritizedDimension = BQVectorUtils.discretize(fieldInfo.getVectorDimension(), 64);
451-
BinaryQuantizer quantizer = new BinaryQuantizer(
436+
final long vectorDataOffset = binarizedVectorData.alignFilePointer(Float.BYTES);
437+
final int descritizedDimension = BQVectorUtils.discretize(fieldInfo.getVectorDimension(), 64);
438+
final BinaryQuantizer quantizer = new BinaryQuantizer(
452439
fieldInfo.getVectorDimension(),
453440
descritizedDimension,
454441
fieldInfo.getVectorSimilarityFunction()
455442
);
443+
444+
IndexOutput tempQuantizedVectorData = null;
445+
IndexOutput tempScoreQuantizedVectorData = null;
446+
final DocsWithFieldSet docsWithField;
447+
boolean success = false;
448+
456449
try {
450+
tempQuantizedVectorData = segmentWriteState.directory.createTempOutput(
451+
binarizedVectorData.getName(),
452+
"temp",
453+
segmentWriteState.context
454+
);
455+
tempScoreQuantizedVectorData = segmentWriteState.directory.createTempOutput(
456+
binarizedVectorData.getName(),
457+
"score_temp",
458+
segmentWriteState.context
459+
);
457460
FloatVectorValues floatVectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
458461
if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
459462
floatVectorValues = new NormalizedFloatVectorValues(floatVectorValues);
460463
}
461-
DocsWithFieldSet docsWithField = writeBinarizedVectorAndQueryData(
464+
docsWithField = writeBinarizedVectorAndQueryData(
462465
tempQuantizedVectorData,
463466
tempScoreQuantizedVectorData,
464467
floatVectorValues,
465468
centroid,
466469
quantizer
467470
);
468471
CodecUtil.writeFooter(tempQuantizedVectorData);
469-
IOUtils.close(tempQuantizedVectorData);
472+
CodecUtil.writeFooter(tempScoreQuantizedVectorData);
473+
success = true;
474+
} finally {
475+
if (success) {
476+
IOUtils.close(tempQuantizedVectorData, tempScoreQuantizedVectorData);
477+
} else {
478+
IOUtils.closeWhileHandlingException(tempQuantizedVectorData, tempScoreQuantizedVectorData);
479+
if (tempQuantizedVectorData != null) {
480+
IOUtils.deleteFilesIgnoringExceptions(segmentWriteState.directory, tempQuantizedVectorData.getName());
481+
}
482+
if (tempScoreQuantizedVectorData != null) {
483+
IOUtils.deleteFilesIgnoringExceptions(segmentWriteState.directory, tempScoreQuantizedVectorData.getName());
484+
}
485+
}
486+
}
487+
488+
IndexInput binarizedDataInput = null;
489+
IndexInput binarizedScoreDataInput = null;
490+
success = false;
491+
try {
470492
binarizedDataInput = segmentWriteState.directory.openInput(tempQuantizedVectorData.getName(), segmentWriteState.context);
471493
binarizedVectorData.copyBytes(binarizedDataInput, binarizedDataInput.length() - CodecUtil.footerLength());
472-
long vectorDataLength = binarizedVectorData.getFilePointer() - vectorDataOffset;
494+
final long vectorDataLength = binarizedVectorData.getFilePointer() - vectorDataOffset;
473495
CodecUtil.retrieveChecksum(binarizedDataInput);
474-
CodecUtil.writeFooter(tempScoreQuantizedVectorData);
475-
IOUtils.close(tempScoreQuantizedVectorData);
476496
binarizedScoreDataInput = segmentWriteState.directory.openInput(
477497
tempScoreQuantizedVectorData.getName(),
478498
segmentWriteState.context
@@ -486,10 +506,9 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
486506
cDotC,
487507
docsWithField
488508
);
489-
success = true;
490509
final IndexInput finalBinarizedDataInput = binarizedDataInput;
491510
final IndexInput finalBinarizedScoreDataInput = binarizedScoreDataInput;
492-
OffHeapBinarizedVectorValues vectorValues = new OffHeapBinarizedVectorValues.DenseOffHeapVectorValues(
511+
final OffHeapBinarizedVectorValues vectorValues = new OffHeapBinarizedVectorValues.DenseOffHeapVectorValues(
493512
fieldInfo.getVectorDimension(),
494513
docsWithField.cardinality(),
495514
centroid,
@@ -499,7 +518,7 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
499518
vectorsScorer,
500519
finalBinarizedDataInput
501520
);
502-
RandomVectorScorerSupplier scorerSupplier = vectorsScorer.getRandomVectorScorerSupplier(
521+
final RandomVectorScorerSupplier scorerSupplier = vectorsScorer.getRandomVectorScorerSupplier(
503522
fieldInfo.getVectorSimilarityFunction(),
504523
new OffHeapBinarizedQueryVectorValues(
505524
finalBinarizedScoreDataInput,
@@ -509,22 +528,20 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
509528
),
510529
vectorValues
511530
);
531+
final String tempQuantizedVectorDataName = tempQuantizedVectorData.getName();
532+
final String tempScoreQuantizedVectorDataName = tempScoreQuantizedVectorData.getName();
533+
success = true;
512534
return new BinarizedCloseableRandomVectorScorerSupplier(scorerSupplier, vectorValues, () -> {
513535
IOUtils.close(finalBinarizedDataInput, finalBinarizedScoreDataInput);
514536
IOUtils.deleteFilesIgnoringExceptions(
515537
segmentWriteState.directory,
516-
tempQuantizedVectorData.getName(),
517-
tempScoreQuantizedVectorData.getName()
538+
tempQuantizedVectorDataName,
539+
tempScoreQuantizedVectorDataName
518540
);
519541
});
520542
} finally {
521543
if (success == false) {
522-
IOUtils.closeWhileHandlingException(
523-
tempQuantizedVectorData,
524-
tempScoreQuantizedVectorData,
525-
binarizedDataInput,
526-
binarizedScoreDataInput
527-
);
544+
IOUtils.closeWhileHandlingException(binarizedDataInput, binarizedScoreDataInput);
528545
IOUtils.deleteFilesIgnoringExceptions(
529546
segmentWriteState.directory,
530547
tempQuantizedVectorData.getName(),

0 commit comments

Comments
 (0)