feat: using "interceptor" to enhance the api(compress)

src/main/java/com/xiaomi/infra/pegasus/client/PegasusClient.java

@@ -43,17 +43,19 @@ public long hash(byte[] key) {
   }
 
   private PegasusTable getTable(String tableName) throws PException {
-    return getTable(tableName, 0);
+    return getTable(tableName, 0, false);
   }
 
-  private PegasusTable getTable(String tableName, int backupRequestDelayMs) throws PException {
+  private PegasusTable getTable(
+      String tableName, int backupRequestDelayMs, boolean enableCompression) throws PException {
     PegasusTable table = tableMap.get(tableName);
     if (table == null) {
       synchronized (tableMapLock) {
         table = tableMap.get(tableName);
         if (table == null) {
           try {
-            TableOptions options = new TableOptions(new PegasusHasher(), backupRequestDelayMs);
+            TableOptions options =
+                new TableOptions(new PegasusHasher(), backupRequestDelayMs, enableCompression);
             Table internalTable = cluster.openTable(tableName, options);
             table = new PegasusTable(this, internalTable);
           } catch (Throwable e) {
@@ -189,9 +191,9 @@ public PegasusTableInterface openTable(String tableName) throws PException {
   }
 
   @Override
-  public PegasusTableInterface openTable(String tableName, int backupRequestDelayMs)
-      throws PException {
-    return getTable(tableName, backupRequestDelayMs);
+  public PegasusTableInterface openTable(
+      String tableName, int backupRequestDelayMs, boolean enableCompression) throws PException {
+    return getTable(tableName, backupRequestDelayMs, enableCompression);
   }
 
   @Override

src/main/java/com/xiaomi/infra/pegasus/client/PegasusClientInterface.java

@@ -58,11 +58,12 @@ public interface PegasusClientInterface {
    *     administrator
    * @param backupRequestDelayMs the delay time to send backup request. If backupRequestDelayMs <=
    *     0, The backup request is disabled.
+   * @param enableCompression whether to enable the table data open auto-compress
    * @return the table handler
    * @throws PException throws exception if any error occurs.
    */
-  public PegasusTableInterface openTable(String tableName, int backupRequestDelayMs)
-      throws PException;
+  public PegasusTableInterface openTable(
+      String tableName, int backupRequestDelayMs, boolean enableCompression) throws PException;
 
   /**
    * Check value exist by key from the cluster

src/main/java/com/xiaomi/infra/pegasus/operator/rrdb_check_and_mutate_operator.java

@@ -47,6 +47,10 @@ public check_and_mutate_response get_response() {
     return resp;
   }
 
+  public check_and_mutate_request get_request() {
+    return request;
+  }
+
   private check_and_mutate_request request;
   private check_and_mutate_response resp;
 }

src/main/java/com/xiaomi/infra/pegasus/operator/rrdb_check_and_set_operator.java

@@ -47,6 +47,10 @@ public check_and_set_response get_response() {
     return resp;
   }
 
+  public check_and_set_request get_request() {
+    return request;
+  }
+
   private check_and_set_request request;
   private check_and_set_response resp;
 }

src/main/java/com/xiaomi/infra/pegasus/operator/rrdb_multi_put_operator.java

@@ -47,6 +47,10 @@ public update_response get_response() {
     return resp;
   }
 
+  public multi_put_request get_request() {
+    return request;
+  }
+
   private multi_put_request request;
   private update_response resp;
 }

src/main/java/com/xiaomi/infra/pegasus/operator/rrdb_put_operator.java

@@ -46,6 +46,10 @@ public update_response get_response() {
     return resp;
   }
 
+  public update_request get_request() {
+    return request;
+  }
+
   private update_request request;
   private update_response resp;
 }

src/main/java/com/xiaomi/infra/pegasus/rpc/TableOptions.java

@@ -7,6 +7,7 @@
 public class TableOptions {
   private final KeyHasher keyHasher;
   private final int backupRequestDelayMs;
+  private final boolean enableCompression;
 
   public KeyHasher keyHasher() {
     return this.keyHasher;
@@ -17,15 +18,20 @@ public int backupRequestDelayMs() {
   }
 
   public static TableOptions forTest() {
-    return new TableOptions(KeyHasher.DEFAULT, 0);
+    return new TableOptions(KeyHasher.DEFAULT, 0, false);
   }
 
-  public TableOptions(KeyHasher h, int backupRequestDelay) {
+  public TableOptions(KeyHasher h, int backupRequestDelay, boolean enableCompression) {
     this.keyHasher = h;
     this.backupRequestDelayMs = backupRequestDelay;
+    this.enableCompression = enableCompression;
   }
 
   public boolean enableBackupRequest() {
     return backupRequestDelayMs > 0;
   }
+
+  public boolean enableCompression() {
+    return enableCompression;
+  }
 }

src/main/java/com/xiaomi/infra/pegasus/rpc/interceptor/BackupRequestInterceptor.java

@@ -11,12 +11,6 @@
 
 public class BackupRequestInterceptor implements TableInterceptor {
 
-  private boolean isOpen;
-
-  public BackupRequestInterceptor(boolean isOpen) {
-    this.isOpen = isOpen;
-  }
-
   @Override
   public void before(ClientRequestRound clientRequestRound, TableHandler tableHandler) {
     backupCall(clientRequestRound, tableHandler);
@@ -33,7 +27,7 @@ public void after(
   }
 
   private void backupCall(ClientRequestRound clientRequestRound, TableHandler tableHandler) {
-    if (!isOpen || !clientRequestRound.getOperator().supportBackupRequest()) {
+    if (!clientRequestRound.getOperator().supportBackupRequest()) {
       return;
     }
 

src/main/java/com/xiaomi/infra/pegasus/rpc/interceptor/CompressionInterceptor.java

@@ -0,0 +1,90 @@
+package com.xiaomi.infra.pegasus.rpc.interceptor;
+
+import com.xiaomi.infra.pegasus.apps.key_value;
+import com.xiaomi.infra.pegasus.apps.mutate;
+import com.xiaomi.infra.pegasus.base.error_code.error_types;
+import com.xiaomi.infra.pegasus.operator.client_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_check_and_mutate_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_check_and_set_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_get_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_multi_get_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_multi_put_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_put_operator;
+import com.xiaomi.infra.pegasus.operator.rrdb_scan_operator;
+import com.xiaomi.infra.pegasus.rpc.async.ClientRequestRound;
+import com.xiaomi.infra.pegasus.rpc.async.TableHandler;
+import com.xiaomi.infra.pegasus.tools.ZstdWrapper;
+import java.util.List;
+
+public class CompressionInterceptor implements TableInterceptor {
+
+  @Override
+  public void before(ClientRequestRound clientRequestRound, TableHandler tableHandler) {
+    tryCompress(clientRequestRound);
+  }
+
+  @Override
+  public void after(
+      ClientRequestRound clientRequestRound, error_types errno, TableHandler tableHandler) {
+    if (errno != error_types.ERR_OK) {
+      return;
+    }
+    tryDecompress(clientRequestRound);
+  }
+
+  private void tryCompress(ClientRequestRound clientRequestRound) {
+    client_operator operator = clientRequestRound.getOperator();
+    if (operator instanceof rrdb_put_operator) {
+      rrdb_put_operator put = (rrdb_put_operator) operator;
+      put.get_request().value.data = ZstdWrapper.compress(put.get_request().value.data);
+      return;
+    }
+
+    if (operator instanceof rrdb_multi_put_operator) {
+      List<key_value> kvs = ((rrdb_multi_put_operator) operator).get_request().kvs;
+      for (key_value kv : kvs) {
+        kv.value.data = ZstdWrapper.compress(kv.value.data);
+      }
+      return;
+    }
+
+    if (operator instanceof rrdb_check_and_set_operator) {
+      rrdb_check_and_set_operator check_and_set = (rrdb_check_and_set_operator) operator;
+      check_and_set.get_request().set_value.data =
+          ZstdWrapper.compress(check_and_set.get_request().set_value.data);
+      return;
+    }
+
+    if (operator instanceof rrdb_check_and_mutate_operator) {
+      List<mutate> mutates = ((rrdb_check_and_mutate_operator) operator).get_request().mutate_list;
+      for (mutate mu : mutates) {
+        mu.value.data = ZstdWrapper.compress(mu.value.data);
+      }
+    }
+  }
+
+  private void tryDecompress(ClientRequestRound clientRequestRound) {
+    client_operator operator = clientRequestRound.getOperator();
+
+    if (operator instanceof rrdb_get_operator) {
+      rrdb_get_operator get = (rrdb_get_operator) operator;
+      get.get_response().value.data = ZstdWrapper.tryDecompress(get.get_response().value.data);
+      return;
+    }
+
+    if (operator instanceof rrdb_multi_get_operator) {
+      List<key_value> kvs = ((rrdb_multi_get_operator) operator).get_response().kvs;
+      for (key_value kv : kvs) {
+        kv.value.data = ZstdWrapper.tryDecompress(kv.value.data);
+      }
+      return;
+    }
+
+    if (operator instanceof rrdb_scan_operator) {
+      List<key_value> kvs = ((rrdb_scan_operator) operator).get_response().kvs;
+      for (key_value kv : kvs) {
+        kv.value.data = ZstdWrapper.tryDecompress(kv.value.data);
+      }
+    }
+  }
+}

src/main/java/com/xiaomi/infra/pegasus/rpc/interceptor/InterceptorManger.java

@@ -12,11 +12,14 @@ public class InterceptorManger {
   private List<TableInterceptor> interceptors = new ArrayList<>();
 
   public InterceptorManger(TableOptions options) {
-    register(new BackupRequestInterceptor(options.enableBackupRequest()));
+    this.register(new BackupRequestInterceptor(), options.enableBackupRequest())
+        .register(new CompressionInterceptor(), options.enableCompression());
   }
 
-  private InterceptorManger register(TableInterceptor interceptor) {
-    interceptors.add(interceptor);
+  private InterceptorManger register(TableInterceptor interceptor, boolean enable) {
+    if (enable) {
+      interceptors.add(interceptor);
+    }
     return this;
   }
 

src/main/java/com/xiaomi/infra/pegasus/tools/ZstdWrapper.java

@@ -12,6 +12,23 @@ public class ZstdWrapper {
 
   private ZstdWrapper() {}
 
+  /**
+   * try decompress the `src`, return `src` directly if failed
+   *
+   * @param src
+   * @return
+   */
+  public static byte[] tryDecompress(byte[] src) {
+    byte[] decompressedValue;
+    try {
+      decompressedValue = decompress(src);
+    } catch (PException e) {
+      // decompress fail
+      decompressedValue = src;
+    }
+    return decompressedValue;
+  }
+
   /**
    * Compresses the `src` and returns the compressed.
    *

src/test/java/com/xiaomi/infra/pegasus/rpc/async/InterceptorTest.java

@@ -0,0 +1,37 @@
+package com.xiaomi.infra.pegasus.rpc.async;
+
+import com.xiaomi.infra.pegasus.client.ClientOptions;
+import com.xiaomi.infra.pegasus.client.PException;
+import com.xiaomi.infra.pegasus.client.PegasusClientFactory;
+import com.xiaomi.infra.pegasus.client.PegasusTableInterface;
+import org.junit.Test;
+import org.junit.jupiter.api.Assertions;
+
+public class InterceptorTest {
+  @Test
+  public void testCompressionInterceptor() throws PException {
+    PegasusTableInterface commonTable =
+        PegasusClientFactory.createClient(ClientOptions.create()).openTable("temp", 0, false);
+    PegasusTableInterface compressTable =
+        PegasusClientFactory.createClient(ClientOptions.create()).openTable("temp", 0, true);
+
+    byte[] hashKey = "hashKey".getBytes();
+    byte[] sortKey = "sortKey".getBytes();
+    byte[] commonValue = "commonValue".getBytes();
+    byte[] compressionValue = "compressionValue".getBytes();
+
+    // if origin value was not compressed, both commonTable and compressTable can read origin value
+    commonTable.set(hashKey, sortKey, commonValue, 10000);
+    Assertions.assertEquals(
+        new String(commonTable.get(hashKey, sortKey, 10000)), new String(commonValue));
+    Assertions.assertEquals(
+        new String(compressTable.get(hashKey, sortKey, 10000)), new String(commonValue));
+
+    // if origin value was compressed, only compressTable can read successfully
+    compressTable.set(hashKey, sortKey, compressionValue, 10000);
+    Assertions.assertNotEquals(
+        new String(commonTable.get(hashKey, sortKey, 10000)), new String(compressionValue));
+    Assertions.assertEquals(
+        new String(compressTable.get(hashKey, sortKey, 10000)), new String(compressionValue));
+  }
+}