Skip to content

Commit

Permalink
Integrate zstd with security manager
Browse files Browse the repository at this point in the history
  • Loading branch information
@dnhatn
dnhatn committed Jan 26, 2024
1 parent 7d69cf9 commit 9eee71b
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 44 deletions.
78 changes: 40 additions & 38 deletions server/src/main/java/org/elasticsearch/index/codec/postings/ES814InlineFieldsConsumer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,6 @@

package org.elasticsearch.index.codec.postings;

import com.github.luben.zstd.ZstdOutputStreamNoFinalizer;

import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.FieldsConsumer;
import org.apache.lucene.codecs.NormsProducer;
Expand All @@ -24,6 +22,7 @@
import org.apache.lucene.store.DataOutput;
import org.apache.lucene.store.IndexOutput;
import org.apache.lucene.util.BytesRef;
import org.apache.lucene.util.BytesRefBuilder;
import org.apache.lucene.util.FixedBitSet;
import org.apache.lucene.util.packed.MonotonicBlockPackedWriter;
import org.apache.lucene.util.packed.PackedInts;
Expand All @@ -32,7 +31,6 @@
import org.elasticsearch.lucene.util.BitStreamOutput;

import java.io.IOException;
import java.io.OutputStream;

import static org.elasticsearch.index.codec.postings.ES814InlinePostingsFormat.INVERTED_INDEX_CODEC;
import static org.elasticsearch.index.codec.postings.ES814InlinePostingsFormat.INVERTED_INDEX_EXTENSION;
Expand Down Expand Up @@ -119,9 +117,8 @@ public void write(Fields fields, NormsProducer norms) throws IOException {
flags |= PostingsEnum.OFFSETS;
}

ByteBuffersDataOutput termsOut = new ByteBuffersDataOutput();
ZstdDataOutput termsOut = new ZstdDataOutput();
ByteBuffersDataOutput postingsOut = new ByteBuffersDataOutput();
ByteBuffersDataOutput sparse = new ByteBuffersDataOutput();

PostingsWriter writer = new PostingsWriter(hasFreqs, hasPositions, hasOffsets, hasPayloads);
TermIndexWriter termIndexWriter = new TermIndexWriter(termIndex, index.getFilePointer());
Expand Down Expand Up @@ -152,16 +149,16 @@ public void write(Fields fields, NormsProducer norms) throws IOException {
++numPending;
maxTermLength = Math.max(maxTermLength, term.length);
++numTerms;
if (numPending >= MAX_BLOCK_TERMS || (termsOut.size() + postingsOut.size()) >= MAX_BLOCK_BYTES) {
if (numPending >= MAX_BLOCK_TERMS || (termsOut.rawSize() + postingsOut.size()) >= MAX_BLOCK_BYTES) {
termIndexWriter.addTerm(term, index.getFilePointer());
flushBlock(numPending, termsOut, postingsOut, sparse);
flushBlock(numPending, termsOut, postingsOut);
numPending = 0;
++numBlocks;
}
}
termIndexWriter.finish(index.getFilePointer());
if (numPending > 0) {
flushBlock(numPending, termsOut, postingsOut, sparse);
flushBlock(numPending, termsOut, postingsOut);
++numBlocks;
}
meta.writeLong(numTerms);
Expand All @@ -188,18 +185,15 @@ public void write(Fields fields, NormsProducer norms) throws IOException {
}
}

private void flushBlock(int numPending, ByteBuffersDataOutput termsOut, ByteBuffersDataOutput postingsOut, ByteBuffersDataOutput sparse)
throws IOException {
private void flushBlock(int numPending, ZstdDataOutput termsOut, ByteBuffersDataOutput postingsOut) throws IOException {
index.writeVInt(numPending); // num terms
compressTerms(termsOut, sparse);
index.writeVLong(termsOut.size()); // original term size
index.writeVLong(sparse.size()); // compressed terms byte
index.writeVLong(termsOut.rawSize()); // original term size
termsOut.compress();
index.writeVLong(termsOut.compressedSize()); // compressed terms byte
index.writeVLong(postingsOut.size()); // postings bytes
sparse.copyTo(index);
termsOut.copyTo(index);
postingsOut.copyTo(index);
postingsOut.reset();
termsOut.reset();
sparse.reset();
}

private class PostingsWriter {
Expand Down Expand Up @@ -414,30 +408,38 @@ void finish(long blockAddress) throws IOException {
}
}

static void compressTerms(ByteBuffersDataOutput in, ByteBuffersDataOutput out) throws IOException {
ZstdOutputStreamNoFinalizer zstream = new ZstdOutputStreamNoFinalizer(new OutputStream() {
@Override
public void write(int b) {
out.writeByte((byte) b);
}
static final class ZstdDataOutput extends DataOutput {
private final BytesRefBuilder raw = new BytesRefBuilder();
private final BytesRefBuilder compressed = new BytesRefBuilder();

@Override
public void write(byte[] b, int off, int len) {
out.writeBytes(b, off, len);
}
}, ZSTD_COMPRESSION_LEVEL);
in.copyTo(new DataOutput() {
@Override
public void writeByte(byte b) throws IOException {
zstream.write(b);
}
@Override
public void writeByte(byte b) {
raw.append(b);
}

@Override
public void writeBytes(byte[] b, int offset, int length) throws IOException {
zstream.write(b, offset, length);
}
});
zstream.close();
@Override
public void writeBytes(byte[] b, int offset, int length) throws IOException {
raw.append(b, offset, length);
}

void compress() {
compressed.setLength(0);
Zstd.compress(raw, compressed, ZSTD_COMPRESSION_LEVEL);
raw.setLength(0);
}

long rawSize() {
return raw.length();
}

long compressedSize() {
return compressed.length();
}

void copyTo(DataOutput out) throws IOException {
out.writeBytes(compressed.bytes(), 0, compressed.length());
compressed.setLength(0);
}
}

@Override
Expand Down
8 changes: 2 additions & 6 deletions server/src/main/java/org/elasticsearch/index/codec/postings/ES814InlineFieldsProducer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,8 +8,6 @@

package org.elasticsearch.index.codec.postings;

import com.github.luben.zstd.Zstd;

import org.apache.lucene.codecs.CodecUtil;
import org.apache.lucene.codecs.FieldsProducer;
import org.apache.lucene.index.BaseTermsEnum;
Expand Down Expand Up @@ -457,13 +455,11 @@ public BytesRef next() throws IOException {
}

private void decompressTerms(int compressedBytes, int originalBytes) throws IOException {
// TODO: Why ZstdInputStreamNoFinalizer is so slow?
zCompressedTerms = ArrayUtil.growNoCopy(zCompressedTerms, compressedBytes);
zDecompressedTerms = ArrayUtil.growNoCopy(zDecompressedTerms, originalBytes);
index.readBytes(zCompressedTerms, 0, compressedBytes);
long actual = Zstd.decompressByteArray(zDecompressedTerms, 0, originalBytes, zCompressedTerms, 0, compressedBytes);
assert actual == originalBytes : actual + " != " + originalBytes;
termsReader = new ByteArrayDataInput(zDecompressedTerms, 0, (int) actual);
Zstd.decompress(zDecompressedTerms, originalBytes, zCompressedTerms, compressedBytes);
termsReader = new ByteArrayDataInput(zDecompressedTerms, 0, originalBytes);
}

private void loadFrame() throws IOException {
Expand Down
48 changes: 48 additions & 0 deletions server/src/main/java/org/elasticsearch/index/codec/postings/Zstd.java
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
/*
* Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
* or more contributor license agreements. Licensed under the Elastic License
* 2.0 and the Server Side Public License, v 1; you may not use this file except
* in compliance with, at your election, the Elastic License 2.0 or the Server
* Side Public License, v 1.
*/

package org.elasticsearch.index.codec.postings;

import org.apache.lucene.util.BytesRefBuilder;

import java.security.AccessController;
import java.security.PrivilegedAction;

public final class Zstd {

public static void compress(BytesRefBuilder raw, BytesRefBuilder out, int compressionLevel) {
AccessController.doPrivileged(new PrivilegedAction<Void>() {
@Override
public Void run() {
int estimatedBytes = Math.toIntExact(com.github.luben.zstd.Zstd.compressBound(raw.length()));
out.grow(estimatedBytes);
long len = com.github.luben.zstd.Zstd.compressByteArray(
out.bytes(),
0,
estimatedBytes,
raw.bytes(),
0,
raw.length(),
compressionLevel
);
out.setLength(Math.toIntExact(len));
return null;
}
});
}

public static void decompress(byte[] dst, int dstSize, byte[] src, int srcSize) {
AccessController.doPrivileged(new PrivilegedAction<Void>() {
@Override
public Void run() {
com.github.luben.zstd.Zstd.decompressByteArray(dst, 0, dstSize, src, 0, srcSize);
return null;
}
});
}
}
3 changes: 3 additions & 0 deletions server/src/main/resources/org/elasticsearch/bootstrap/security.policy
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,9 @@ grant codeBase "${codebase.elasticsearch}" {

// for plugin api dynamic settings instances
permission java.lang.RuntimePermission "accessClassInPackage.jdk.internal.reflect";

// zstd - should isolate this permission in a separate module
permission java.lang.RuntimePermission "loadLibrary.*";
};

//// Very special jar permissions:
Expand Down

0 comments on commit 9eee71b

Please sign in to comment.