diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoBytesField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoBytesField.java index d4c41f9..104a140 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoBytesField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoBytesField.java @@ -157,6 +157,18 @@ public void serialize(PrintWriter w) { w.format("_i += _%sLen;\n", ccName); } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("%s;\n", writeTagToBufExpr(tagName())); + w.format("LightProtoCodec.writeVarInt(_b, _%sLen);\n", ccName); + w.format("if (_%sIdx == -1) {\n", ccName); + // Absolute-indexed copy for the same reason as the array path. + w.format(" %s.getBytes(%s.readerIndex(), _b, _%sLen);\n", ccName, ccName, ccName); + w.format("} else {\n"); + w.format(" _parsedBuffer.getBytes(_%sIdx, _b, _%sLen);\n", ccName, ccName); + w.format("}\n"); + } + @Override public void materialize(PrintWriter w) { diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java index 5f0e843..765e737 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoField.java @@ -155,6 +155,13 @@ public void fieldClear(PrintWriter w, String enclosingType) { abstract public void serialize(PrintWriter w); + /** + * Emit code that writes this field through the ByteBuf API ({@code _b}): the + * allocation-free path used for messages too large to stage in a scratch array. + * Must produce bytes identical to {@link #serialize(PrintWriter)}. + */ + abstract public void serializeToBuf(PrintWriter w); + abstract public void serializeJson(PrintWriter w); abstract public void parseJson(PrintWriter w); @@ -220,6 +227,14 @@ protected String writeTagExpr(String tag) { } } + protected String writeTagToBufExpr(String tag) { + if (field.getNumber() <= 15) { + return String.format("_b.writeByte(%s)", tag); + } else { + return String.format("LightProtoCodec.writeVarInt(_b, %s)", tag); + } + } + protected String tagName() { return "_" + Util.upperCase(field.getName(), "tag"); } diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java index 95da500..88acdb0 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMapField.java @@ -858,6 +858,36 @@ private void generateValueDataSize(PrintWriter w, String idxVar) { } } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("for (int _entryIdx = 0; _entryIdx < _%sCount; _entryIdx++) {\n", ccName); + + // Compute entry size + w.format(" int _entrySize = 0;\n"); + + // Key size: 1 (tag) + data size + w.format(" _entrySize += 1;\n"); // key tag is always 1 byte + generateKeyDataSize(w, "_entryIdx"); + + // Value size: 1 (tag) + data size + w.format(" _entrySize += 1;\n"); // value tag is always 1 byte + generateValueDataSize(w, "_entryIdx"); + + // Write outer tag + entry size + w.format(" %s;\n", writeTagToBufExpr(tagName())); + w.format(" LightProtoCodec.writeVarInt(_b, _entrySize);\n"); + + // Write key tag + key data + w.format(" _b.writeByte(%s);\n", keyTagConstant()); + generateSerializeKeyDataToBuf(w, "_entryIdx"); + + // Write value tag + value data + w.format(" _b.writeByte(%s);\n", valueTagConstant()); + generateSerializeValueDataToBuf(w, "_entryIdx"); + + w.format("}\n"); + } + private void generateSerializeKeyData(PrintWriter w, String idxVar) { if (isStringKey()) { w.format(" LightProtoCodec.StringHolder _ksh = _%sKeys[%s];\n", ccName, idxVar); @@ -901,6 +931,46 @@ private void generateSerializeValueData(PrintWriter w, String idxVar) { } } + private void generateSerializeKeyDataToBuf(PrintWriter w, String idxVar) { + if (isStringKey()) { + w.format(" LightProtoCodec.StringHolder _ksh = _%sKeys[%s];\n", ccName, idxVar); + w.format(" LightProtoCodec.writeVarInt(_b, _ksh.len);\n"); + w.format(" if (_ksh.idx == -1) {\n"); + w.format(" LightProtoCodec.writeString(_b, _ksh.s, _ksh.len);\n"); + w.format(" } else {\n"); + w.format(" _parsedBuffer.getBytes(_ksh.idx, _b, _ksh.len);\n"); + w.format(" }\n"); + } else { + LightProtoNumberField.serializeNumberToBuf(w, keyField, String.format("_%sKeys[%s]", ccName, idxVar)); + } + } + + private void generateSerializeValueDataToBuf(PrintWriter w, String idxVar) { + if (isStringValue()) { + w.format(" LightProtoCodec.StringHolder _vsh = _%sValues[%s];\n", ccName, idxVar); + w.format(" LightProtoCodec.writeVarInt(_b, _vsh.len);\n"); + w.format(" if (_vsh.idx == -1) {\n"); + w.format(" LightProtoCodec.writeString(_b, _vsh.s, _vsh.len);\n"); + w.format(" } else {\n"); + w.format(" _parsedBuffer.getBytes(_vsh.idx, _b, _vsh.len);\n"); + w.format(" }\n"); + } else if (isBytesValue()) { + w.format(" LightProtoCodec.BytesHolder _vbh = _%sValues[%s];\n", ccName, idxVar); + w.format(" LightProtoCodec.writeVarInt(_b, _vbh.len);\n"); + w.format(" if (_vbh.idx == -1) {\n"); + w.format(" _vbh.b.getBytes(_vbh.b.readerIndex(), _b, _vbh.len);\n"); + w.format(" } else {\n"); + w.format(" _parsedBuffer.getBytes(_vbh.idx, _b, _vbh.len);\n"); + w.format(" }\n"); + } else if (isMessageValue()) { + w.format(" LightProtoCodec.writeVarInt(_b, _%sValues[%s].getSerializedSize());\n", + ccName, idxVar); + w.format(" _%sValues[%s]._writeTo(_b);\n", ccName, idxVar); + } else { + LightProtoNumberField.serializeNumberToBuf(w, valueField, String.format("_%sValues[%s]", ccName, idxVar)); + } + } + @Override public void serializedSize(PrintWriter w) { w.format("for (int _i = 0; _i < _%sCount; _i++) {\n", ccName); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java index 1829821..1e74f1a 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessage.java @@ -376,16 +376,23 @@ private void generateSerialize(PrintWriter w) { w.format(" int _writeIdx = _b.writerIndex();\n"); w.format(" _writeTo(_b.array(), _b.arrayOffset() + _writeIdx);\n"); w.format(" _b.writerIndex(_writeIdx + _serializedSize);\n"); + w.format(" } else if (_serializedSize > LightProtoCodec.SCRATCH_RETAIN_MAX) {\n"); + // Messages too large for a retained scratch array write through the + // ByteBuf API field by field: slower per element, but allocation-free. + // Staging them would allocate a fresh full-size heap array on every + // write (multi-MB arrays are G1 humongous allocations), invisible to + // any direct-memory accounting sized to the target buffer. + w.format(" _b.ensureWritable(_serializedSize);\n"); + w.format(" _writeTo(_b);\n"); w.format(" } else {\n"); // Direct, composite and other buffers: compose in a scratch array cached // on this (typically pooled) instance and transfer with a single bulk // write. Plain byte[] stores compile to raw memory accesses on every JDK, // unlike sun.misc.Unsafe accesses which carry a per-call deprecation - // check since JDK 24. + // check since JDK 24. The dispatch above bounds _serializedSize by + // SCRATCH_RETAIN_MAX, so the scratch is always retainable. w.format(" byte[] _s = LightProtoCodec.scratchFor(this._scratch, _serializedSize);\n"); - w.format(" if (_s.length <= LightProtoCodec.SCRATCH_RETAIN_MAX) {\n"); - w.format(" this._scratch = _s;\n"); - w.format(" }\n"); + w.format(" this._scratch = _s;\n"); w.format(" _writeTo(_s, 0);\n"); w.format(" _b.writeBytes(_s, 0, _serializedSize);\n"); w.format(" }\n"); @@ -423,6 +430,37 @@ private void generateSerialize(PrintWriter w) { w.format(" return _i;\n"); w.format(" }\n"); + + w.println(" /**"); + w.println(" * Internal: serialize this message field by field through the ByteBuf"); + w.println(" * API. The allocation-free path for messages larger than"); + w.println(" * {@code SCRATCH_RETAIN_MAX}; nested messages write through as well, so"); + w.println(" * no element of the tree stages in a scratch array. Public only so that"); + w.println(" * generated messages in other packages can serialize nested fields of"); + w.println(" * this type into the same buffer."); + w.println(" */"); + w.format(" public void _writeTo(io.netty.buffer.ByteBuf _b) {\n"); + if (hasRequiredFields()) { + w.format(" checkRequiredFields();\n"); + } + if (useBitDrivenTraversal()) { + // Same set-bit traversal as the array path: field order — and bytes — + // must match it exactly. + emitBitDrivenTraversal(w, f -> f.serializeToBuf(w)); + } else { + for (LightProtoField f : fields) { + String condition = f.serializeCondition(); + if (condition != null) { + w.format(" if (%s) {\n", condition); + f.serializeToBuf(w); + w.format(" }\n"); + } else { + f.serializeToBuf(w); + } + } + } + + w.format(" }\n"); } private void generateGetSerializedSize(PrintWriter w) { diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessageField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessageField.java index 218711b..5b633e5 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessageField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoMessageField.java @@ -104,6 +104,15 @@ public void serialize(PrintWriter w) { w.format("_i = %s._writeTo(_a, _i);\n", ccName); } + @Override + public void serializeToBuf(PrintWriter w) { + // Nested messages write through as well, so no element of the tree + // stages in a scratch array. + w.format("%s;\n", writeTagToBufExpr(tagName())); + w.format("LightProtoCodec.writeVarInt(_b, %s.getSerializedSize());\n", ccName); + w.format("%s._writeTo(_b);\n", ccName); + } + @Override public void clear(PrintWriter w) { w.format("if (%s()){\n", Util.camelCase("has", ccName)); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoNumberField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoNumberField.java index f88a091..04dd64d 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoNumberField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoNumberField.java @@ -80,6 +80,40 @@ static void serializeNumber(PrintWriter w, ProtoFieldDescriptor field, String na } } + static void serializeNumberToBuf(PrintWriter w, ProtoFieldDescriptor field, String name) { + if (field.isEnumField()) { + w.format(" LightProtoCodec.writeVarInt(_b, %s.getValue());\n", name); + } else if (field.getProtoType().equals("bool")) { + w.format(" _b.writeByte(%s ? 1 : 0);\n", name); + } else if (field.getProtoType().equals("int32")) { + w.format(" LightProtoCodec.writeVarInt(_b, %s);\n", name); + } else if (field.getProtoType().equals("uint32")) { + w.format(" LightProtoCodec.writeVarInt(_b, %s);\n", name); + } else if (field.getProtoType().equals("sint32")) { + w.format(" LightProtoCodec.writeSignedVarInt(_b, %s);\n", name); + } else if (field.getProtoType().equals("sint64")) { + w.format(" LightProtoCodec.writeSignedVarInt64(_b, %s);\n", name); + } else if (field.getProtoType().equals("int64")) { + w.format(" LightProtoCodec.writeVarInt64(_b, %s);\n", name); + } else if (field.getProtoType().equals("uint64")) { + w.format(" LightProtoCodec.writeVarInt64(_b, %s);\n", name); + } else if (field.getProtoType().equals("fixed32")) { + w.format(" LightProtoCodec.writeFixedInt32(_b, %s);\n", name); + } else if (field.getProtoType().equals("fixed64")) { + w.format(" LightProtoCodec.writeFixedInt64(_b, %s);\n", name); + } else if (field.getProtoType().equals("sfixed32")) { + w.format(" LightProtoCodec.writeFixedInt32(_b, %s);\n", name); + } else if (field.getProtoType().equals("sfixed64")) { + w.format(" LightProtoCodec.writeFixedInt64(_b, %s);\n", name); + } else if (field.getProtoType().equals("double")) { + w.format(" LightProtoCodec.writeDouble(_b, %s);\n", name); + } else if (field.getProtoType().equals("float")) { + w.format(" LightProtoCodec.writeFloat(_b, %s);\n", name); + } else { + throw new IllegalArgumentException("Failed to write serializer for field: " + field.getProtoType()); + } + } + static String parseNumber(ProtoFieldDescriptor field) { if (field.isEnumField()) { return String.format("%s.valueOf(LightProtoCodec.readVarInt(_buffer))", field.getJavaType()); @@ -193,6 +227,12 @@ public void serialize(PrintWriter w) { serializeNumber(w, field, ccName); } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("%s;\n", writeTagToBufExpr(tagName())); + serializeNumberToBuf(w, field, ccName); + } + @Override public void serializeJson(PrintWriter w) { String type = field.getProtoType(); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedBytesField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedBytesField.java index 9700531..780328b 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedBytesField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedBytesField.java @@ -97,6 +97,20 @@ public void serialize(PrintWriter w) { w.format("}\n"); } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("for (int i = 0; i < _%sCount; i++) {\n", pluralName); + w.format(" LightProtoCodec.BytesHolder _bh = %s[i];\n", pluralName); + w.format(" %s;\n", writeTagToBufExpr(tagName())); + w.format(" LightProtoCodec.writeVarInt(_b, _bh.len);\n"); + w.format(" if (_bh.idx == -1) {\n"); + w.format(" _bh.b.getBytes(_bh.b.readerIndex(), _b, _bh.len);\n"); + w.format(" } else {\n"); + w.format(" _parsedBuffer.getBytes(_bh.idx, _b, _bh.len);\n"); + w.format(" }\n"); + w.format("}\n"); + } + @Override public void serializeJson(PrintWriter w) { w.format("_b.writeByte('[');\n"); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedMessageField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedMessageField.java index 16a4049..d3af615 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedMessageField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedMessageField.java @@ -82,6 +82,16 @@ public void serialize(PrintWriter w) { w.format("}\n"); } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("for (int i = 0; i < _%sCount; i++) {\n", pluralName); + w.format(" %s _item = %s[i];\n", field.getJavaType(), pluralName); + w.format(" %s;\n", writeTagToBufExpr(tagName())); + w.format(" LightProtoCodec.writeVarInt(_b, _item.getSerializedSize());\n"); + w.format(" _item._writeTo(_b);\n"); + w.format("}\n"); + } + @Override public void serializeJson(PrintWriter w) { w.format("_b.writeByte('[');\n"); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedNumberField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedNumberField.java index 7521086..a24af9d 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedNumberField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedNumberField.java @@ -108,6 +108,34 @@ public void serialize(PrintWriter w) { } } + @Override + public void serializeToBuf(PrintWriter w) { + int fixedSize = LightProtoNumberField.fixedDataSize(field); + if (field.isPacked()) { + w.format(" %s;\n", writeTagToBufExpr(tagName() + "_PACKED")); + if (fixedSize >= 0) { + w.format(" LightProtoCodec.writeVarInt(_b, _%sCount * %d);\n", pluralName, fixedSize); + } else { + w.format(" int _%sSize = 0;\n", pluralName); + w.format("for (int i = 0; i < _%sCount; i++) {\n", pluralName); + w.format(" %s _item = %s[i];\n", field.getJavaType(), pluralName); + w.format(" _%sSize += %s;\n", pluralName, LightProtoNumberField.serializedSizeOfNumber(field, "_item")); + w.format("}\n"); + w.format(" LightProtoCodec.writeVarInt(_b, _%sSize);\n", pluralName); + } + w.format("for (int i = 0; i < _%sCount; i++) {\n", pluralName); + w.format(" %s _item = %s[i];\n", field.getJavaType(), pluralName); + LightProtoNumberField.serializeNumberToBuf(w, field, "_item"); + w.format("}\n"); + } else { + w.format("for (int i = 0; i < _%sCount; i++) {\n", pluralName); + w.format(" %s _item = %s[i];\n", field.getJavaType(), pluralName); + w.format(" %s;\n", writeTagToBufExpr(tagName())); + LightProtoNumberField.serializeNumberToBuf(w, field, "_item"); + w.format("}\n"); + } + } + @Override public void serializeJson(PrintWriter w) { w.format("_b.writeByte('[');\n"); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedStringField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedStringField.java index a0477b8..9f80184 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedStringField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoRepeatedStringField.java @@ -95,6 +95,20 @@ public void serialize(PrintWriter w) { w.format("}\n"); } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("for (int i = 0; i < _%sCount; i++) {\n", pluralName); + w.format(" LightProtoCodec.StringHolder _sh = %s[i];\n", pluralName); + w.format(" %s;\n", writeTagToBufExpr(tagName())); + w.format(" LightProtoCodec.writeVarInt(_b, _sh.len);\n"); + w.format(" if (_sh.idx == -1) {\n"); + w.format(" LightProtoCodec.writeString(_b, _sh.s, _sh.len);\n"); + w.format(" } else {\n"); + w.format(" _parsedBuffer.getBytes(_sh.idx, _b, _sh.len);\n"); + w.format(" }\n"); + w.format("}\n"); + } + @Override public void serializeJson(PrintWriter w) { w.format("_b.writeByte('[');\n"); diff --git a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoStringField.java b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoStringField.java index ea37d75..65e1612 100644 --- a/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoStringField.java +++ b/code-generator/src/main/java/io/streamnative/lightproto/generator/LightProtoStringField.java @@ -102,6 +102,17 @@ public void serialize(PrintWriter w) { w.format("}\n"); } + @Override + public void serializeToBuf(PrintWriter w) { + w.format("%s;\n", writeTagToBufExpr(tagName())); + w.format("LightProtoCodec.writeVarInt(_b, _%sBufferLen);\n", ccName); + w.format("if (_%sBufferIdx == -1) {\n", ccName); + w.format(" LightProtoCodec.writeString(_b, %s, _%sBufferLen);\n", ccName, ccName); + w.format("} else {\n"); + w.format(" _parsedBuffer.getBytes(_%sBufferIdx, _b, _%sBufferLen);\n", ccName, ccName); + w.format("}\n"); + } + @Override public void serializeJson(PrintWriter w) { w.format("LightProtoCodec.writeJsonString(_b, %s());\n", Util.camelCase("get", field.getName())); diff --git a/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java b/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java index 77eb322..d0623b7 100644 --- a/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java +++ b/code-generator/src/main/resources/io/streamnative/lightproto/generator/LightProtoCodec.java @@ -394,12 +394,15 @@ static void writeString(ByteBuf b, String s, int bytesCount) { // Serialization composes into a plain byte[] with an int cursor: heap buffers // are written in place through their backing array, other buffer types are // composed in a reusable scratch array and transferred with a single bulk - // writeBytes(). Plain array stores compile to raw memory accesses on every JDK - // (no sun.misc.Unsafe in the hot loop — its memory-access methods carry a - // per-call deprecation check since JDK 24). - - // Scratch arrays larger than this are not retained on the message instance, - // so outlier messages don't pin large allocations. + // writeBytes(). Messages larger than SCRATCH_RETAIN_MAX skip the scratch and + // write through the ByteBuf API field by field instead. Plain array stores + // compile to raw memory accesses on every JDK (no sun.misc.Unsafe in the hot + // loop — its memory-access methods carry a per-call deprecation check since + // JDK 24). + + // Messages larger than this serialize field by field through the ByteBuf API + // instead of composing in a scratch array, so outlier messages neither pin + // nor churn large heap allocations. static final int SCRATCH_RETAIN_MAX = 1024 * 1024; /** Returns current if it can hold size bytes, otherwise a larger replacement. */ @@ -408,8 +411,8 @@ static byte[] scratchFor(byte[] current, int size) { return current; } if (size > SCRATCH_RETAIN_MAX) { - // The result won't be retained, so growth amortization is pointless: - // allocate exactly what this outlier message needs. + // Not reached from generated code (outlier messages write through the + // ByteBuf API instead); kept so the size contract holds for any caller. return new byte[size]; } // Double to amortize growth, but never past the retain cap: otherwise diff --git a/tests/src/test/java/io/streamnative/lightproto/tests/LargeWriteThroughTest.java b/tests/src/test/java/io/streamnative/lightproto/tests/LargeWriteThroughTest.java new file mode 100644 index 0000000..9d440d4 --- /dev/null +++ b/tests/src/test/java/io/streamnative/lightproto/tests/LargeWriteThroughTest.java @@ -0,0 +1,294 @@ +/** + * Copyright 2026 StreamNative + * + * 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 io.streamnative.lightproto.tests; + +import static org.junit.jupiter.api.Assertions.assertArrayEquals; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.junit.jupiter.api.Assumptions.assumeTrue; + +import io.netty.buffer.ByteBuf; +import io.netty.buffer.CompositeByteBuf; +import io.netty.buffer.PooledByteBufAllocator; +import io.netty.buffer.Unpooled; +import java.lang.management.ManagementFactory; +import java.lang.reflect.Field; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import org.apache.pulsar.common.api.proto.BaseCommand; +import org.apache.pulsar.common.api.proto.PulsarApi; +import org.junit.jupiter.api.Test; + +/** + * Messages larger than {@code SCRATCH_RETAIN_MAX} must serialize to non-array + * buffers through the ByteBuf API — byte-identical to the scratch path, and + * without allocating a transient full-size heap array per write (the 0.8.0 + * regression that OOMed Pulsar's proxy back-pressure test). + */ +public class LargeWriteThroughTest { + + /** + * ~4.5 MB of repeated strings (8192 names, ~550 chars each), mirroring the + * Pulsar CommandGetTopicsOfNamespaceResponse shape that exposed the + * regression. Every 100th name carries non-ASCII characters so the UTF-8 + * write path is exercised at scale. + */ + private static S largeStrings() { + StringBuilder sb = new StringBuilder(); + for (int i = 0; i < 14; i++) { + sb.append("persistent://public/default/large-write-through-").append(i).append('/'); + } + String base = sb.toString(); + S s = new S().setId("large"); + for (int i = 0; i < 8192; i++) { + s.addName(base + (i % 100 == 0 ? "-λ∞≈-" : "-") + i); + } + return s; + } + + private static byte[] drain(ByteBuf b) { + byte[] out = new byte[b.readableBytes()]; + b.readBytes(out); + return out; + } + + private static byte[] scratchOf(Object msg) throws Exception { + Field f = msg.getClass().getDeclaredField("_scratch"); + f.setAccessible(true); + return (byte[]) f.get(msg); + } + + @Test + public void testLargeRepeatedStringsAllTargets() throws Exception { + S s = largeStrings(); + byte[] expected = s.toByteArray(); + assertTrue(expected.length > LightProtoCodec.SCRATCH_RETAIN_MAX); + + Strings.S.Builder pb = Strings.S.newBuilder().setId("large"); + for (int i = 0; i < s.getNamesCount(); i++) { + pb.addNames(s.getNameAt(i)); + } + assertArrayEquals(expected, pb.build().toByteArray()); + + ByteBuf direct = PooledByteBufAllocator.DEFAULT.directBuffer(expected.length); + try { + assertEquals(expected.length, s.writeTo(direct)); + assertArrayEquals(expected, drain(direct)); + + // Re-serialization must reuse the cached sizes and stay identical + direct.clear(); + s.writeTo(direct); + assertArrayEquals(expected, drain(direct)); + } finally { + direct.release(); + } + + CompositeByteBuf composite = Unpooled.compositeBuffer(); + try { + s.writeTo(composite); + assertArrayEquals(expected, drain(composite)); + } finally { + composite.release(); + } + + // A parsed message writes through the lazy-string passthrough branch + S parsed = new S(); + parsed.parseFrom(expected); + ByteBuf direct2 = PooledByteBufAllocator.DEFAULT.directBuffer(expected.length); + try { + parsed.writeTo(direct2); + assertArrayEquals(expected, drain(direct2)); + } finally { + direct2.release(); + } + } + + @Test + public void testLargeBytesFieldWriteThrough() throws Exception { + byte[] payload = new byte[2 * 1024 * 1024]; + new Random(42).nextBytes(payload); + B b = new B().setPayload(payload); + b.addExtraItem(new byte[]{1, 2, 3}); + b.addExtraItem(new byte[]{4, 5}); + + byte[] expected = b.toByteArray(); + assertTrue(expected.length > LightProtoCodec.SCRATCH_RETAIN_MAX); + + ByteBuf direct = PooledByteBufAllocator.DEFAULT.directBuffer(expected.length); + try { + b.writeTo(direct); + assertArrayEquals(expected, drain(direct)); + + // Parsed passthrough: bytes fields copy straight from the parse buffer + B parsed = new B(); + parsed.parseFrom(expected); + direct.clear(); + parsed.writeTo(direct); + assertArrayEquals(expected, drain(direct)); + } finally { + direct.release(); + } + } + + @Test + public void testLargeNestedMessageTreeWriteThrough() { + M m = new M(); + m.setX().setA("a-value").setB("b-value"); + for (int i = 0; i < 9000; i++) { + M.KV kv = m.addItem(); + kv.setK("key-" + "k".repeat(60) + "-" + i); + kv.setV("val-" + "v".repeat(60) + "-" + i); + if (i % 10 == 0) { + kv.setXx().setN(i); + } + } + + byte[] expected = m.toByteArray(); + assertTrue(expected.length > LightProtoCodec.SCRATCH_RETAIN_MAX); + + ByteBuf direct = PooledByteBufAllocator.DEFAULT.directBuffer(expected.length); + try { + m.writeTo(direct); + assertArrayEquals(expected, drain(direct)); + } finally { + direct.release(); + } + } + + @Test + public void testPulsarBaseCommandShapeWriteThrough() { + // The exact production shape: a BaseCommand wrapping a multi-MB + // CommandGetTopicsOfNamespaceResponse. BaseCommand also exercises the + // bit-driven traversal variant of the write-through path. + List topics = new ArrayList<>(); + String base = "persistent://public/default/" + "t".repeat(520) + "-"; + for (int i = 0; i < 8192; i++) { + topics.add(base + i); + } + + BaseCommand cmd = new BaseCommand().setType(BaseCommand.Type.GET_TOPICS_OF_NAMESPACE_RESPONSE); + cmd.setGetTopicsOfNamespaceResponse().setRequestId(42).addAllTopics(topics); + + PulsarApi.BaseCommand pb = PulsarApi.BaseCommand.newBuilder() + .setType(PulsarApi.BaseCommand.Type.GET_TOPICS_OF_NAMESPACE_RESPONSE) + .setGetTopicsOfNamespaceResponse(PulsarApi.CommandGetTopicsOfNamespaceResponse.newBuilder() + .setRequestId(42) + .addAllTopics(topics)) + .build(); + + byte[] expected = pb.toByteArray(); + assertTrue(expected.length > LightProtoCodec.SCRATCH_RETAIN_MAX); + assertEquals(expected.length, cmd.getSerializedSize()); + + ByteBuf direct = PooledByteBufAllocator.DEFAULT.directBuffer(expected.length); + try { + cmd.writeTo(direct); + assertArrayEquals(expected, drain(direct)); + } finally { + direct.release(); + } + } + + @Test + public void testWriteThroughAllocationBoundBytes() throws Exception { + // Bytes fields transfer with bulk getBytes() in every JVM config, so + // this asserts the core property unconditionally: no transient + // full-size heap array per write. + byte[] payload = new byte[5 * 1024 * 1024]; + new Random(7).nextBytes(payload); + B b = new B().setPayload(payload); + assertAllocationFreeWrites(b, b.getSerializedSize()); + } + + @Test + public void testWriteThroughAllocationBoundStrings() throws Exception { + // With -XX:-CompactStrings both write paths copy each string's bytes + // through a temporary array (writeString/writeRawString cannot read the + // String's internal byte[]), so the string-heavy variant of this + // assertion only holds under the default compact-strings config. + assumeTrue(!ManagementFactory.getRuntimeMXBean().getInputArguments().contains("-XX:-CompactStrings")); + + S s = largeStrings(); + assertAllocationFreeWrites(s, s.getSerializedSize()); + } + + private static void assertAllocationFreeWrites(LightProtoCodec.LightProtoMessage msg, int size) + throws Exception { + var mxBean = ManagementFactory.getThreadMXBean(); + assumeTrue(mxBean instanceof com.sun.management.ThreadMXBean); + com.sun.management.ThreadMXBean tb = (com.sun.management.ThreadMXBean) mxBean; + assumeTrue(tb.isThreadAllocatedMemorySupported()); + if (!tb.isThreadAllocatedMemoryEnabled()) { + tb.setThreadAllocatedMemoryEnabled(true); + } + + assertTrue(size > LightProtoCodec.SCRATCH_RETAIN_MAX); + + ByteBuf direct = PooledByteBufAllocator.DEFAULT.directBuffer(size + 64); + try { + for (int i = 0; i < 3; i++) { + direct.clear(); + msg.writeTo(direct); + } + + long tid = Thread.currentThread().getId(); + long before = tb.getThreadAllocatedBytes(tid); + for (int i = 0; i < 5; i++) { + direct.clear(); + msg.writeTo(direct); + } + long allocated = tb.getThreadAllocatedBytes(tid) - before; + + // Staging would allocate a fresh full-size array per write (5x the + // message size over this loop); the write-through path allocates + // none of that. The bound leaves generous slack for JIT/runtime + // noise. + assertTrue(allocated < size / 4, + "expected allocation-free serialization but " + allocated + + " bytes were allocated for 5 writes of a " + size + "-byte message"); + // The scratch path was never taken for this outlier message + assertNull(scratchOf(msg)); + } finally { + direct.release(); + } + } + + @Test + public void testScratchStillUsedBelowThreshold() throws Exception { + // ~600 KB: stays on the scratch fast path and retains the array + S s = new S().setId("medium"); + for (int i = 0; i < 1200; i++) { + s.addName("x".repeat(500)); + } + byte[] expected = s.toByteArray(); + assertTrue(expected.length <= LightProtoCodec.SCRATCH_RETAIN_MAX); + + ByteBuf direct = PooledByteBufAllocator.DEFAULT.directBuffer(expected.length); + try { + s.writeTo(direct); + assertArrayEquals(expected, drain(direct)); + } finally { + direct.release(); + } + + byte[] scratch = scratchOf(s); + assertNotNull(scratch); + assertTrue(scratch.length <= LightProtoCodec.SCRATCH_RETAIN_MAX); + } +}