Make bytes available

Author: shangm2

drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DataSizeToBytesThriftCodec.java

@@ -42,7 +42,7 @@ public DataSizeToBytesThriftCodec(ThriftCatalog thriftCatalog)
     }
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DurationToMillisThriftCodec.java

@@ -42,7 +42,7 @@ public DurationToMillisThriftCodec(ThriftCatalog thriftCatalog)
     }
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/JodaDateTimeToEpochMillisThriftCodec.java

@@ -41,7 +41,7 @@ public JodaDateTimeToEpochMillisThriftCodec(ThriftCatalog thriftCatalog)
     }
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/LocaleToLanguageTagCodec.java

@@ -42,7 +42,7 @@ public LocaleToLanguageTagCodec(ThriftCatalog thriftCatalog)
     }
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalDoubleThriftCodec.java

@@ -17,6 +17,7 @@
 
 import com.facebook.drift.codec.CodecThriftType;
 import com.facebook.drift.codec.ThriftCodec;
+import com.facebook.drift.codec.metadata.ThriftCatalog;
 import com.facebook.drift.codec.metadata.ThriftType;
 import com.facebook.drift.protocol.TProtocolReader;
 import com.facebook.drift.protocol.TProtocolWriter;
@@ -34,7 +35,7 @@ public class OptionalDoubleThriftCodec
     private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.DOUBLE, OptionalDouble.class, OptionalDouble.empty());
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalIntThriftCodec.java

@@ -17,6 +17,7 @@
 
 import com.facebook.drift.codec.CodecThriftType;
 import com.facebook.drift.codec.ThriftCodec;
+import com.facebook.drift.codec.metadata.ThriftCatalog;
 import com.facebook.drift.codec.metadata.ThriftType;
 import com.facebook.drift.protocol.TProtocolReader;
 import com.facebook.drift.protocol.TProtocolWriter;
@@ -40,7 +41,7 @@ public ThriftType getType()
     }
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalLongThriftCodec.java

@@ -17,6 +17,7 @@
 
 import com.facebook.drift.codec.CodecThriftType;
 import com.facebook.drift.codec.ThriftCodec;
+import com.facebook.drift.codec.metadata.ThriftCatalog;
 import com.facebook.drift.codec.metadata.ThriftType;
 import com.facebook.drift.protocol.TProtocolReader;
 import com.facebook.drift.protocol.TProtocolWriter;
@@ -34,7 +35,7 @@ public class OptionalLongThriftCodec
     private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.I64, OptionalLong.class, OptionalLong.empty());
 
     @CodecThriftType
-    public static ThriftType getThriftType()
+    public static ThriftType getThriftType(ThriftCatalog catalog)
     {
         return THRIFT_TYPE;
     }

drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGenerator.java

@@ -162,16 +162,16 @@ private void loadCustomCodec(String customCodecClassName)
                     throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must be static: " + customCodecClassName + "#" + method.getName());
                 }
 
-                if (method.getParameterCount() != 0) {
-                    throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must have no parameters: " + customCodecClassName + "#" + method.getName());
+                if (method.getParameterCount() != 1 || method.getParameterTypes()[0] != ThriftCatalog.class) {
+                    throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must have only 1 parameter of ThriftCatalog type: " + customCodecClassName + "#" + method.getName());
                 }
 
                 if (!ThriftType.class.isAssignableFrom(method.getReturnType())) {
                     throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must return ThriftType: " + customCodecClassName + "#" + method.getName());
                 }
 
                 try {
-                    thriftType = (ThriftType) method.invoke(null);
+                    thriftType = (ThriftType) method.invoke(null, this.catalog);
                     if (thriftType == null) {
                         throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType returns null: " + customCodecClassName + "#" + method.getName());
                     }

drift-protocol/src/main/java/com/facebook/drift/protocol/TByteArrayOutputStream.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright (C) 2017 Facebook, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.drift.protocol;
+
+import java.io.ByteArrayOutputStream;
+import java.nio.charset.Charset;
+
+public class TByteArrayOutputStream
+        extends ByteArrayOutputStream
+{
+    public TByteArrayOutputStream(int size)
+    {
+        super(size);
+    }
+
+    public TByteArrayOutputStream()
+    {
+        super();
+    }
+
+    protected byte[] get()
+    {
+        return buf;
+    }
+
+    public int len()
+    {
+        return count;
+    }
+
+    public String toString(Charset charset)
+    {
+        return new String(buf, 0, count, charset);
+    }
+}

drift-protocol/src/main/java/com/facebook/drift/protocol/TMemoryBuffer.java

@@ -15,44 +15,70 @@
  */
 package com.facebook.drift.protocol;
 
-import static java.lang.Math.max;
+import java.io.UnsupportedEncodingException;
 
 public class TMemoryBuffer
         implements TTransport
 {
-    private byte[] data;
-    private int head;
-    private int tail;
+    // The contents of the buffer
+    private TByteArrayOutputStream data;
 
-    public TMemoryBuffer(int initialSize)
+    // Position to read next byte from
+    private int position;
+
+    public TMemoryBuffer(int size)
     {
-        data = new byte[max(initialSize, 16)];
+        reset(size);
     }
 
     @Override
     public void read(byte[] buf, int off, int len)
             throws TTransportException
     {
-        if (head - tail < len) {
-            throw new TTransportException("Too few bytes in buffer");
+        byte[] src = data.get();
+        int amtToRead = (len > data.len() - position ? data.len() - position : len);
+        if (amtToRead > 0) {
+            System.arraycopy(src, position, buf, off, amtToRead);
+            position += amtToRead;
         }
-        System.arraycopy(data, tail, buf, off, len);
-        tail += len;
     }
 
     @Override
     public void write(byte[] buf, int off, int len)
     {
-        int available = data.length - head;
-        if (available < len) {
-            int need = len - available - tail;
-            byte[] temp = new byte[max(data.length * 2, need)];
-            System.arraycopy(data, tail, temp, 0, head - tail);
-            data = temp;
-            head -= tail;
-            tail = 0;
-        }
-        System.arraycopy(buf, off, data, head, len);
-        head += len;
+        data.write(buf, off, len);
+    }
+
+    /**
+     * Output the contents of the memory buffer as a String, using the supplied encoding
+     *
+     * @param enc the encoding to use
+     * @return the contents of the memory buffer as a String
+     */
+    public String toString(String enc)
+            throws UnsupportedEncodingException
+    {
+        return data.toString(enc);
+    }
+
+    public int length()
+    {
+        return data.len();
+    }
+
+    public void reset(int size)
+    {
+        data = new TByteArrayOutputStream(size);
+        position = 0;
+    }
+
+    public byte[] getBytes()
+    {
+        return data.get();
+    }
+
+    public int getArrayPos()
+    {
+        return position;
     }
 }