Skip to content

Commit

Permalink
Fix potential file leak in ES816BinaryQuantizedVectorsWriter (elastic…
Browse files Browse the repository at this point in the history
…#120014) (elastic#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
  • Loading branch information
iverase authored Jan 14, 2025
1 parent 7141ff1 commit 0710a9d
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 32 deletions.
6 changes: 6 additions & 0 deletions docs/changelog/120014.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@
pr: 120014
summary: Fix potential file leak in ES816BinaryQuantizedVectorsWriter
area: Search
type: bug
issues:
- 119981
Original file line number Diff line number Diff line change
Expand Up @@ -433,46 +433,66 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
float[] centroid,
float cDotC
) throws IOException {
long vectorDataOffset = binarizedVectorData.alignFilePointer(Float.BYTES);
final IndexOutput tempQuantizedVectorData = segmentWriteState.directory.createTempOutput(
binarizedVectorData.getName(),
"temp",
segmentWriteState.context
);
final IndexOutput tempScoreQuantizedVectorData = segmentWriteState.directory.createTempOutput(
binarizedVectorData.getName(),
"score_temp",
segmentWriteState.context
);
IndexInput binarizedDataInput = null;
IndexInput binarizedScoreDataInput = null;
boolean success = false;
int descritizedDimension = BQVectorUtils.discretize(fieldInfo.getVectorDimension(), 64);
BinaryQuantizer quantizer = new BinaryQuantizer(
final long vectorDataOffset = binarizedVectorData.alignFilePointer(Float.BYTES);
final int descritizedDimension = BQVectorUtils.discretize(fieldInfo.getVectorDimension(), 64);
final BinaryQuantizer quantizer = new BinaryQuantizer(
fieldInfo.getVectorDimension(),
descritizedDimension,
fieldInfo.getVectorSimilarityFunction()
);

IndexOutput tempQuantizedVectorData = null;
IndexOutput tempScoreQuantizedVectorData = null;
final DocsWithFieldSet docsWithField;
boolean success = false;

try {
tempQuantizedVectorData = segmentWriteState.directory.createTempOutput(
binarizedVectorData.getName(),
"temp",
segmentWriteState.context
);
tempScoreQuantizedVectorData = segmentWriteState.directory.createTempOutput(
binarizedVectorData.getName(),
"score_temp",
segmentWriteState.context
);
FloatVectorValues floatVectorValues = KnnVectorsWriter.MergedVectorValues.mergeFloatVectorValues(fieldInfo, mergeState);
if (fieldInfo.getVectorSimilarityFunction() == COSINE) {
floatVectorValues = new NormalizedFloatVectorValues(floatVectorValues);
}
DocsWithFieldSet docsWithField = writeBinarizedVectorAndQueryData(
docsWithField = writeBinarizedVectorAndQueryData(
tempQuantizedVectorData,
tempScoreQuantizedVectorData,
floatVectorValues,
centroid,
quantizer
);
CodecUtil.writeFooter(tempQuantizedVectorData);
IOUtils.close(tempQuantizedVectorData);
CodecUtil.writeFooter(tempScoreQuantizedVectorData);
success = true;
} finally {
if (success) {
IOUtils.close(tempQuantizedVectorData, tempScoreQuantizedVectorData);
} else {
IOUtils.closeWhileHandlingException(tempQuantizedVectorData, tempScoreQuantizedVectorData);
if (tempQuantizedVectorData != null) {
IOUtils.deleteFilesIgnoringExceptions(segmentWriteState.directory, tempQuantizedVectorData.getName());
}
if (tempScoreQuantizedVectorData != null) {
IOUtils.deleteFilesIgnoringExceptions(segmentWriteState.directory, tempScoreQuantizedVectorData.getName());
}
}
}

IndexInput binarizedDataInput = null;
IndexInput binarizedScoreDataInput = null;
success = false;
try {
binarizedDataInput = segmentWriteState.directory.openInput(tempQuantizedVectorData.getName(), segmentWriteState.context);
binarizedVectorData.copyBytes(binarizedDataInput, binarizedDataInput.length() - CodecUtil.footerLength());
long vectorDataLength = binarizedVectorData.getFilePointer() - vectorDataOffset;
final long vectorDataLength = binarizedVectorData.getFilePointer() - vectorDataOffset;
CodecUtil.retrieveChecksum(binarizedDataInput);
CodecUtil.writeFooter(tempScoreQuantizedVectorData);
IOUtils.close(tempScoreQuantizedVectorData);
binarizedScoreDataInput = segmentWriteState.directory.openInput(
tempScoreQuantizedVectorData.getName(),
segmentWriteState.context
Expand All @@ -486,10 +506,9 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
cDotC,
docsWithField
);
success = true;
final IndexInput finalBinarizedDataInput = binarizedDataInput;
final IndexInput finalBinarizedScoreDataInput = binarizedScoreDataInput;
OffHeapBinarizedVectorValues vectorValues = new OffHeapBinarizedVectorValues.DenseOffHeapVectorValues(
final OffHeapBinarizedVectorValues vectorValues = new OffHeapBinarizedVectorValues.DenseOffHeapVectorValues(
fieldInfo.getVectorDimension(),
docsWithField.cardinality(),
centroid,
Expand All @@ -499,7 +518,7 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
vectorsScorer,
finalBinarizedDataInput
);
RandomVectorScorerSupplier scorerSupplier = vectorsScorer.getRandomVectorScorerSupplier(
final RandomVectorScorerSupplier scorerSupplier = vectorsScorer.getRandomVectorScorerSupplier(
fieldInfo.getVectorSimilarityFunction(),
new OffHeapBinarizedQueryVectorValues(
finalBinarizedScoreDataInput,
Expand All @@ -509,22 +528,20 @@ private CloseableRandomVectorScorerSupplier mergeOneFieldToIndex(
),
vectorValues
);
final String tempQuantizedVectorDataName = tempQuantizedVectorData.getName();
final String tempScoreQuantizedVectorDataName = tempScoreQuantizedVectorData.getName();
success = true;
return new BinarizedCloseableRandomVectorScorerSupplier(scorerSupplier, vectorValues, () -> {
IOUtils.close(finalBinarizedDataInput, finalBinarizedScoreDataInput);
IOUtils.deleteFilesIgnoringExceptions(
segmentWriteState.directory,
tempQuantizedVectorData.getName(),
tempScoreQuantizedVectorData.getName()
tempQuantizedVectorDataName,
tempScoreQuantizedVectorDataName
);
});
} finally {
if (success == false) {
IOUtils.closeWhileHandlingException(
tempQuantizedVectorData,
tempScoreQuantizedVectorData,
binarizedDataInput,
binarizedScoreDataInput
);
IOUtils.closeWhileHandlingException(binarizedDataInput, binarizedScoreDataInput);
IOUtils.deleteFilesIgnoringExceptions(
segmentWriteState.directory,
tempQuantizedVectorData.getName(),
Expand Down

0 comments on commit 0710a9d

Please sign in to comment.