-
Notifications
You must be signed in to change notification settings - Fork 86
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[client-common] Added safeguard for compressor (#1307)
* [client-common] Added safeguard for compressor Today, the `compress`/`decompress` can still be invoked even the compressor is closed already and for zstd based compressor, it would crash. This PR add some safeguard and fail fast if the compressor is already closed. * Fixed integration test failures * Minor tweak * Added a unit test * Fixed minor comment * Skipped locking for NoopCompressor
- Loading branch information
Showing
5 changed files
with
143 additions
and
22 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
74 changes: 67 additions & 7 deletions
74
.../venice-client-common/src/main/java/com/linkedin/venice/compression/VeniceCompressor.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,43 +1,103 @@ | ||
package com.linkedin.venice.compression; | ||
|
||
import com.linkedin.venice.exceptions.VeniceException; | ||
import com.linkedin.venice.utils.ByteUtils; | ||
import java.io.Closeable; | ||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.nio.ByteBuffer; | ||
import java.util.concurrent.locks.ReentrantReadWriteLock; | ||
|
||
|
||
public abstract class VeniceCompressor implements Closeable { | ||
protected static final int SCHEMA_HEADER_LENGTH = ByteUtils.SIZE_OF_INT; | ||
private final CompressionStrategy compressionStrategy; | ||
private boolean isClosed = false; | ||
/** | ||
* To avoid the race condition between 'compress'/'decompress' operation and 'close'. | ||
*/ | ||
private final ReentrantReadWriteLock readWriteLock = new ReentrantReadWriteLock(); | ||
|
||
protected VeniceCompressor(CompressionStrategy compressionStrategy) { | ||
this.compressionStrategy = compressionStrategy; | ||
} | ||
|
||
public abstract byte[] compress(byte[] data) throws IOException; | ||
interface CompressionRunnable<R> { | ||
R run() throws IOException; | ||
} | ||
|
||
private <R> R executeWithSafeGuard(CompressionRunnable<R> runnable) throws IOException { | ||
readWriteLock.readLock().lock(); | ||
try { | ||
if (isClosed) { | ||
throw new VeniceException("Compressor for " + getCompressionStrategy() + " has been closed"); | ||
} | ||
return runnable.run(); | ||
} finally { | ||
readWriteLock.readLock().unlock(); | ||
} | ||
} | ||
|
||
public byte[] compress(byte[] data) throws IOException { | ||
return executeWithSafeGuard(() -> compressInternal(data)); | ||
} | ||
|
||
public abstract ByteBuffer compress(ByteBuffer src, int startPositionOfOutput) throws IOException; | ||
protected abstract byte[] compressInternal(byte[] data) throws IOException; | ||
|
||
public abstract ByteBuffer decompress(ByteBuffer data) throws IOException; | ||
public ByteBuffer compress(ByteBuffer src, int startPositionOfOutput) throws IOException { | ||
return executeWithSafeGuard(() -> compressInternal(src, startPositionOfOutput)); | ||
} | ||
|
||
public abstract ByteBuffer decompress(byte[] data, int offset, int length) throws IOException; | ||
protected abstract ByteBuffer compressInternal(ByteBuffer src, int startPositionOfOutput) throws IOException; | ||
|
||
public ByteBuffer decompress(ByteBuffer data) throws IOException { | ||
return executeWithSafeGuard(() -> decompressInternal(data)); | ||
} | ||
|
||
protected abstract ByteBuffer decompressInternal(ByteBuffer data) throws IOException; | ||
|
||
public ByteBuffer decompress(byte[] data, int offset, int length) throws IOException { | ||
return executeWithSafeGuard(() -> decompressInternal(data, offset, length)); | ||
} | ||
|
||
protected abstract ByteBuffer decompressInternal(byte[] data, int offset, int length) throws IOException; | ||
|
||
/** | ||
* This method tries to decompress data and maybe prepend the schema header. | ||
* The returned ByteBuffer will be backed by byte array that starts with schema header, followed by the | ||
* decompressed data. The ByteBuffer will be positioned at the beginning of the decompressed data and the remaining of | ||
* the ByteBuffer will be the length of the decompressed data. | ||
*/ | ||
public abstract ByteBuffer decompressAndPrependSchemaHeader(byte[] data, int offset, int length, int schemaHeader) | ||
throws IOException; | ||
public ByteBuffer decompressAndPrependSchemaHeader(byte[] data, int offset, int length, int schemaHeader) | ||
throws IOException { | ||
return executeWithSafeGuard(() -> decompressAndPrependSchemaHeaderInternal(data, offset, length, schemaHeader)); | ||
} | ||
|
||
protected abstract ByteBuffer decompressAndPrependSchemaHeaderInternal( | ||
byte[] data, | ||
int offset, | ||
int length, | ||
int schemaHeader) throws IOException; | ||
|
||
public CompressionStrategy getCompressionStrategy() { | ||
return compressionStrategy; | ||
} | ||
|
||
public abstract InputStream decompress(InputStream inputStream) throws IOException; | ||
public InputStream decompress(InputStream inputStream) throws IOException { | ||
return executeWithSafeGuard(() -> decompressInternal(inputStream)); | ||
} | ||
|
||
protected abstract InputStream decompressInternal(InputStream inputStream) throws IOException; | ||
|
||
public void close() throws IOException { | ||
readWriteLock.writeLock().lock(); | ||
try { | ||
isClosed = true; | ||
closeInternal(); | ||
} finally { | ||
readWriteLock.writeLock().unlock(); | ||
} | ||
} | ||
|
||
protected abstract void closeInternal() throws IOException; | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters