Make bytes available

shangm2 · shangm2 · commit 3d204e5e811c · 2025-06-04T12:36:46.000-07:00
diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DataSizeToBytesThriftCodec.java b/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;
     }
diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/DurationToMillisThriftCodec.java b/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;
     }
diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/JodaDateTimeToEpochMillisThriftCodec.java b/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;
     }
diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/LocaleToLanguageTagCodec.java b/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;
     }
diff --git a/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/UuidToLeachSalzBinaryEncodingThriftCodec.java b/drift-codec-utils/src/main/java/com/facebook/drift/codec/utils/UuidToLeachSalzBinaryEncodingThriftCodec.java
@@ -15,6 +15,7 @@
  */
 package com.facebook.drift.codec.utils;
 
+import com.facebook.drift.codec.CodecThriftType;
 import com.facebook.drift.codec.ThriftCodec;
 import com.facebook.drift.codec.internal.coercion.FromThrift;
 import com.facebook.drift.codec.internal.coercion.ToThrift;
@@ -33,24 +34,32 @@
 
 /**
  * This codec will encode UUIDs into the format specified in RFC 4122:
- *
+ * <p>
  * https://www.ietf.org/rfc/rfc4122.txt
- *
+ * <p>
  * Likewise, it presumes the input is of this format.
  */
 public class UuidToLeachSalzBinaryEncodingThriftCodec
         implements ThriftCodec<UUID>
 {
+    private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.BINARY, UUID.class);
+
     @Inject
     public UuidToLeachSalzBinaryEncodingThriftCodec(ThriftCatalog thriftCatalog)
     {
         thriftCatalog.addDefaultCoercions(getClass());
     }
 
+    @CodecThriftType
+    public static ThriftType getThriftType(ThriftCatalog catalog)
+    {
+        return THRIFT_TYPE;
+    }
+
     @Override
     public ThriftType getType()
     {
-        return new ThriftType(ThriftType.BINARY, UUID.class);
+        return THRIFT_TYPE;
     }
 
     @Override
diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalDoubleThriftCodec.java b/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;
     }
diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalIntThriftCodec.java b/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;
     }
diff --git a/drift-codec/src/main/java/com/facebook/drift/codec/internal/builtin/OptionalLongThriftCodec.java b/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;
     }
diff --git a/drift-idl-generator/src/main/java/com/facebook/drift/idl/generator/ThriftIdlGenerator.java b/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());
                     }
diff --git a/drift-protocol/src/main/java/com/facebook/drift/protocol/TByteArrayOutputStream.java b/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);
+    }
+}
diff --git a/drift-protocol/src/main/java/com/facebook/drift/protocol/TMemoryBuffer.java b/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;
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ public DataSizeToBytesThriftCodec(ThriftCatalog thriftCatalog)`
`42`	`42`	`}`
`43`	`43`
`44`	`44`	`@CodecThriftType`
`45`		`- public static ThriftType getThriftType()`
	`45`	`+ public static ThriftType getThriftType(ThriftCatalog catalog)`
`46`	`46`	`{`
`47`	`47`	`return THRIFT_TYPE;`
`48`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ public DurationToMillisThriftCodec(ThriftCatalog thriftCatalog)`
`42`	`42`	`}`
`43`	`43`
`44`	`44`	`@CodecThriftType`
`45`		`- public static ThriftType getThriftType()`
	`45`	`+ public static ThriftType getThriftType(ThriftCatalog catalog)`
`46`	`46`	`{`
`47`	`47`	`return THRIFT_TYPE;`
`48`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -41,7 +41,7 @@ public JodaDateTimeToEpochMillisThriftCodec(ThriftCatalog thriftCatalog)`
`41`	`41`	`}`
`42`	`42`
`43`	`43`	`@CodecThriftType`
`44`		`- public static ThriftType getThriftType()`
	`44`	`+ public static ThriftType getThriftType(ThriftCatalog catalog)`
`45`	`45`	`{`
`46`	`46`	`return THRIFT_TYPE;`
`47`	`47`	`}`
Original file line number	Diff line number	Diff line change
`@@ -42,7 +42,7 @@ public LocaleToLanguageTagCodec(ThriftCatalog thriftCatalog)`
`42`	`42`	`}`
`43`	`43`
`44`	`44`	`@CodecThriftType`
`45`		`- public static ThriftType getThriftType()`
	`45`	`+ public static ThriftType getThriftType(ThriftCatalog catalog)`
`46`	`46`	`{`
`47`	`47`	`return THRIFT_TYPE;`
`48`	`48`	`}`
Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@`
`17`	`17`
`18`	`18`	`import com.facebook.drift.codec.CodecThriftType;`
`19`	`19`	`import com.facebook.drift.codec.ThriftCodec;`
	`20`	`+import com.facebook.drift.codec.metadata.ThriftCatalog;`
`20`	`21`	`import com.facebook.drift.codec.metadata.ThriftType;`
`21`	`22`	`import com.facebook.drift.protocol.TProtocolReader;`
`22`	`23`	`import com.facebook.drift.protocol.TProtocolWriter;`
`@@ -34,7 +35,7 @@ public class OptionalDoubleThriftCodec`
`34`	`35`	`private static final ThriftType THRIFT_TYPE = new ThriftType(ThriftType.DOUBLE, OptionalDouble.class, OptionalDouble.empty());`
`35`	`36`
`36`	`37`	`@CodecThriftType`
`37`		`- public static ThriftType getThriftType()`
	`38`	`+ public static ThriftType getThriftType(ThriftCatalog catalog)`
`38`	`39`	`{`
`39`	`40`	`return THRIFT_TYPE;`
`40`	`41`	`}`
Original file line number	Diff line number	Diff line change
`@@ -162,16 +162,16 @@ private void loadCustomCodec(String customCodecClassName)`
`162`	`162`	`throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must be static: " + customCodecClassName + "#" + method.getName());`
`163`	`163`	`}`
`164`	`164`
`165`		`- if (method.getParameterCount() != 0) {`
`166`		`- throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must have no parameters: " + customCodecClassName + "#" + method.getName());`
	`165`	`+ if (method.getParameterCount() != 1 \|\| method.getParameterTypes()[0] != ThriftCatalog.class) {`
	`166`	`+ throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must have only 1 parameter of ThriftCatalog type: " + customCodecClassName + "#" + method.getName());`
`167`	`167`	`}`
`168`	`168`
`169`	`169`	`if (!ThriftType.class.isAssignableFrom(method.getReturnType())) {`
`170`	`170`	`throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType must return ThriftType: " + customCodecClassName + "#" + method.getName());`
`171`	`171`	`}`
`172`	`172`
`173`	`173`	`try {`
`174`		`- thriftType = (ThriftType) method.invoke(null);`
	`174`	`+ thriftType = (ThriftType) method.invoke(null, this.catalog);`
`175`	`175`	`if (thriftType == null) {`
`176`	`176`	`throw new ThriftIdlGeneratorException("Method annotated with @CodecThriftType returns null: " + customCodecClassName + "#" + method.getName());`
`177`	`177`	`}`