HBASE-28613 Use streaming when marshalling protobuf REST output (#5943)

Signed-off-by: Ankit Singhal <[email protected]>
(cherry picked from commit d1d8b4d)

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/ProtobufMessageHandler.java

@@ -17,7 +17,10 @@
  */
 package org.apache.hadoop.hbase.rest;
 
+import com.google.protobuf.CodedOutputStream;
+import com.google.protobuf.Message;
 import java.io.IOException;
+import java.io.OutputStream;
 import org.apache.yetus.audience.InterfaceAudience;
 
 /**
@@ -26,13 +29,48 @@
  */
 @InterfaceAudience.Private
 public interface ProtobufMessageHandler {
-  /** Returns the protobuf represention of the model */
-  byte[] createProtobufOutput();
+
+  // The Jetty 9.4 HttpOutput default commit size is 32K/4 = 8K. We use that size to avoid
+  // double buffering (and copying) in HttpOutput. If we ever increase the HttpOutput commit size,
+  // we need to adjust this accordingly. We should also revisit this when Jetty is upgraded.
+  int BUFFER_SIZE = 8 * 1024;
+
+  /** Writes the protobuf represention of the model to os */
+  default void writeProtobufOutput(OutputStream os) throws IOException {
+    // Creating an explicit CodedOutputStream for the following reasons :
+    // 1. This avoids the cost of pre-computing the message size
+    // 2. This lets us set the buffer size explicitly
+    CodedOutputStream cos = CodedOutputStream.newInstance(os, BUFFER_SIZE);
+    messageFromObject().writeTo(cos);
+    cos.flush();
+  }
+
+  /**
+   * Returns the protobuf represention of the model in a byte array Use
+   * {@link org.apache.hadoop.hbase.rest.ProtobufMessageHandler#writeProtobufOutput(OutputStream)}
+   * for better performance
+   * @return the protobuf encoded object in a byte array
+   */
+  default byte[] createProtobufOutput() {
+    return messageFromObject().toByteArray();
+  }
+
+  /**
+   * Convert to model to a protobuf Message object
+   * @return the protobuf Message object
+   */
+  Message messageFromObject();
 
   /**
    * Initialize the model from a protobuf representation.
    * @param message the raw bytes of the protobuf message
    * @return reference to self for convenience
    */
+  // TODO implement proper stream handling for unmarshalling.
+  // Using byte array here lets us use ProtobufUtil.mergeFrom in the implementations to
+  // avoid the CodedOutputStream size limitation, but is slow
+  // and memory intensive. We could use the ProtobufUtil.mergeFrom() variant that takes
+  // an inputStream and sets the size limit to maxInt.
+  // This would help both on the client side, and when processing large Puts on the server.
   ProtobufMessageHandler getObjectFromMessage(byte[] message) throws IOException;
 }

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/CellModel.java

@@ -21,6 +21,7 @@
 
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import javax.xml.bind.annotation.XmlAccessType;
@@ -200,7 +201,7 @@ public int getValueLength() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     Cell.Builder builder = Cell.newBuilder();
     builder.setColumn(ByteStringer.wrap(getColumn()));
     if (valueLength == MAGIC_LENGTH) {
@@ -211,7 +212,7 @@ public byte[] createProtobufOutput() {
     if (hasUserTimestamp()) {
       builder.setTimestamp(getTimestamp());
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/CellSetModel.java

@@ -19,6 +19,7 @@
 
 import static org.apache.hadoop.hbase.rest.model.CellModel.MAGIC_LENGTH;
 
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -106,7 +107,7 @@ public List<RowModel> getRows() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     CellSet.Builder builder = CellSet.newBuilder();
     for (RowModel row : getRows()) {
       CellSet.Row.Builder rowBuilder = CellSet.Row.newBuilder();
@@ -132,7 +133,7 @@ public byte[] createProtobufOutput() {
       }
       builder.addRows(rowBuilder);
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/NamespacesInstanceModel.java

@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.rest.model;
 
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.HashMap;
@@ -139,7 +140,7 @@ public String toString() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     NamespaceProperties.Builder builder = NamespaceProperties.newBuilder();
     if (properties != null) {
       for (Map.Entry<String, String> entry : properties.entrySet()) {
@@ -150,7 +151,7 @@ public byte[] createProtobufOutput() {
         builder.addProps(property);
       }
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/NamespacesModel.java

@@ -18,6 +18,7 @@
 package org.apache.hadoop.hbase.rest.model;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -94,10 +95,10 @@ public String toString() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     Namespaces.Builder builder = Namespaces.newBuilder();
     builder.addAllNamespace(namespaces);
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/RowModel.java

@@ -20,6 +20,7 @@
 import static org.apache.hadoop.hbase.rest.model.CellModel.MAGIC_LENGTH;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -179,7 +180,7 @@ public List<CellModel> getCells() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     // there is no standalone row protobuf message
     throw new UnsupportedOperationException("no protobuf equivalent to RowModel");
   }

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/ScannerModel.java

@@ -19,6 +19,7 @@
 
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -791,7 +792,7 @@ public void setFilter(String filter) {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     Scanner.Builder builder = Scanner.newBuilder();
     if (!Bytes.equals(startRow, HConstants.EMPTY_START_ROW)) {
       builder.setStartRow(ByteStringer.wrap(startRow));
@@ -821,7 +822,7 @@ public byte[] createProtobufOutput() {
         builder.addLabels(label);
     }
     builder.setCacheBlocks(cacheBlocks);
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/StorageClusterStatusModel.java

@@ -18,6 +18,7 @@
 package org.apache.hadoop.hbase.rest.model;
 
 import com.fasterxml.jackson.annotation.JsonProperty;
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -651,7 +652,7 @@ public String toString() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     StorageClusterStatus.Builder builder = StorageClusterStatus.newBuilder();
     builder.setRegions(regions);
     builder.setRequests(requests);
@@ -686,7 +687,7 @@ public byte[] createProtobufOutput() {
     for (String node : deadNodes) {
       builder.addDeadNodes(node);
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/TableInfoModel.java

@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.rest.model;
 
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -121,7 +122,7 @@ public String toString() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     TableInfo.Builder builder = TableInfo.newBuilder();
     builder.setName(name);
     for (TableRegionModel aRegion : regions) {
@@ -133,7 +134,7 @@ public byte[] createProtobufOutput() {
       regionBuilder.setLocation(aRegion.getLocation());
       builder.addRegions(regionBuilder);
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/TableListModel.java

@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.rest.model;
 
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -89,12 +90,12 @@ public String toString() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     TableList.Builder builder = TableList.newBuilder();
     for (TableModel aTable : tables) {
       builder.addName(aTable.getName());
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/TableSchemaModel.java

@@ -20,6 +20,7 @@
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import java.util.ArrayList;
@@ -245,7 +246,7 @@ public void __setReadOnly(boolean value) {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     TableSchema.Builder builder = TableSchema.newBuilder();
     builder.setName(name);
     for (Map.Entry<QName, Object> e : attrs.entrySet()) {
@@ -278,7 +279,7 @@ public byte[] createProtobufOutput() {
     if (attrs.containsKey(READONLY)) {
       builder.setReadOnly(Boolean.parseBoolean(attrs.get(READONLY).toString()));
     }
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/model/VersionModel.java

@@ -17,6 +17,7 @@
  */
 package org.apache.hadoop.hbase.rest.model;
 
+import com.google.protobuf.Message;
 import java.io.IOException;
 import java.io.Serializable;
 import javax.servlet.ServletContext;
@@ -161,14 +162,14 @@ public String toString() {
   }
 
   @Override
-  public byte[] createProtobufOutput() {
+  public Message messageFromObject() {
     Version.Builder builder = Version.newBuilder();
     builder.setRestVersion(restVersion);
     builder.setJvmVersion(jvmVersion);
     builder.setOsVersion(osVersion);
     builder.setServerVersion(serverVersion);
     builder.setJerseyVersion(jerseyVersion);
-    return builder.build().toByteArray();
+    return builder.build();
   }
 
   @Override

hbase-rest/src/main/java/org/apache/hadoop/hbase/rest/provider/producer/ProtobufMessageBodyProducer.java

@@ -59,6 +59,6 @@ public long getSize(ProtobufMessageHandler m, Class<?> type, Type genericType,
   public void writeTo(ProtobufMessageHandler m, Class<?> type, Type genericType,
     Annotation[] annotations, MediaType mediaType, MultivaluedMap<String, Object> httpHeaders,
     OutputStream entityStream) throws IOException, WebApplicationException {
-    entityStream.write(m.createProtobufOutput());
+    m.writeProtobufOutput(entityStream);
   }
 }