Skip to content

Optimize XML Codec #1231

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
adwsingh merged 3 commits into from
Jun 10, 2026
Merged

Optimize XML Codec #1231

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
2d5a14c
Optimize XML Codec
adwsingh Jun 7, 2026
1d88f1e
Run protocol tests againt multiple codec providers
adwsingh Jun 10, 2026
d8e0956
Address PR comments
adwsingh Jun 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 2 additions & 5 deletions aws/aws-sigv4/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@ import com.google.gradle.osdetector.OsDetector

plugins {
id("smithy-java.module-conventions")
alias(libs.plugins.jmh)
id("smithy-java.jmh-conventions")
alias(libs.plugins.osdetector)
}

Expand Down Expand Up @@ -31,9 +31,6 @@ afterEvaluate {
}

jmh {
iterations = 3
warmupIterations = 2
fork = 1
// profilers.add("async:output=flamegraph")
// profilers.add("gc")
iterations = 3
}
5 changes: 5 additions & 0 deletions aws/client/aws-client-awsjson/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,11 @@ dependencies {
testImplementation(libs.smithy.aws.protocol.tests)
}

protocolTestRuns {
run("native") { systemProperty("smithy-java.json-provider", "smithy") }
run("jackson") { systemProperty("smithy-java.json-provider", "jackson") }
}

val generator = "software.amazon.smithy.java.protocoltests.generators.ProtocolTestGenerator"
addGenerateSrcsTask(generator, "awsJson1_0", "aws.protocoltests.json10#JsonRpc10")
addGenerateSrcsTask(generator, "awsJson1_1", "aws.protocoltests.json#JsonProtocol")
6 changes: 6 additions & 0 deletions aws/client/aws-client-awsquery/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,19 @@ extra["moduleName"] = "software.amazon.smithy.java.aws.client.awsquery"
dependencies {
api(project(":client:client-http"))
api(project(":codecs:xml-codec"))
api(project(":codecs:codec-commons", configuration = "shadow"))
api(project(":io"))
api(libs.smithy.aws.traits)

// Protocol test dependencies
testImplementation(libs.smithy.aws.protocol.tests)
}

protocolTestRuns {
run("native") { systemProperty("smithy-java.xml-provider", "smithy") }
run("stax") { }
}

val generator = "software.amazon.smithy.java.protocoltests.generators.ProtocolTestGenerator"
addGenerateSrcsTask(generator, "awsQuery", "aws.protocoltests.query#AwsQuery")
addGenerateSrcsTask(generator, "ec2Query", "aws.protocoltests.ec2#AwsEc2")
2 changes: 1 addition & 1 deletion ...src/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQueryClientProtocol.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ public <I extends SerializableStruct, O extends SerializableStruct> HttpRequest
SmithyUri endpoint
) {
String operationName = operation.schema().id().getName();
QueryFormSerializer serializer = new QueryFormSerializer(
QueryFormSerializer serializer = QueryFormSerializer.acquire(
QueryFormSerializer.QueryVariant.AWS_QUERY,
operationName,
version);
Expand Down
161 changes: 147 additions & 14 deletions ...c/main/java/software/amazon/smithy/java/aws/client/awsquery/AwsQuerySchemaExtensions.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,31 +5,54 @@

package software.amazon.smithy.java.aws.client.awsquery;

import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
import software.amazon.smithy.aws.traits.protocols.Ec2QueryNameTrait;
import software.amazon.smithy.java.core.schema.Schema;
import software.amazon.smithy.java.core.schema.SchemaExtensionKey;
import software.amazon.smithy.java.core.schema.SchemaExtensionProvider;
import software.amazon.smithy.java.core.schema.TraitKey;
import software.amazon.smithy.model.shapes.ShapeType;
import software.amazon.smithy.model.traits.TimestampFormatTrait;
import software.amazon.smithy.utils.SmithyInternalApi;
import software.amazon.smithy.utils.StringUtils;

/**
* Pre-computes the URL-encoded member-name bytes for AWS Query and EC2 Query protocols once per {@link Schema}.
* Pre-computes URL-encoded member-name bytes and list/map metadata for AWS Query and EC2 Query protocols.
*/
@SmithyInternalApi
public final class AwsQuerySchemaExtensions
implements SchemaExtensionProvider<AwsQuerySchemaExtensions.QueryMemberBinding> {

public static final SchemaExtensionKey<QueryMemberBinding> KEY = new SchemaExtensionKey<>();

private static final byte[] MEMBER_BYTES = "member".getBytes(StandardCharsets.UTF_8);
private static final byte[] KEY_BYTES = "key".getBytes(StandardCharsets.UTF_8);
private static final byte[] VALUE_BYTES = "value".getBytes(StandardCharsets.UTF_8);
private static final byte[] ENTRY_BYTES = "entry".getBytes(StandardCharsets.UTF_8);

/**
* Pre-encoded member-name bytes for both query variants.
* Pre-computed query binding data for a schema.
*
* @param awsQueryNameBytes Bytes to use as the awsQuery name. Never null.
* @param ec2QueryNameBytes Bytes to use as the ec2Query name. Never null.
* @param awsQueryNameBytes Bytes for awsQuery member name. Null for non-members.
* @param ec2QueryNameBytes Bytes for ec2Query member name. Null for non-members.
* @param listFlattened Whether this list member has xmlFlattened trait.
* @param listMemberNameBytes Pre-computed list member name bytes (for non-flattened lists). Null if flattened.
* @param mapFlattened Whether this map member has xmlFlattened trait.
* @param mapKeyNameBytes Pre-computed map key name bytes.
* @param mapValueNameBytes Pre-computed map value name bytes.
* @param mapEntryNameBytes Pre-computed map entry name bytes (null if flattened).
* @param timestampFormat Pre-resolved timestamp format for timestamp members.
*/
public record QueryMemberBinding(byte[] awsQueryNameBytes, byte[] ec2QueryNameBytes) {}
public record QueryMemberBinding(
byte[] awsQueryNameBytes,
byte[] ec2QueryNameBytes,
boolean listFlattened,
byte[] listMemberNameBytes,
boolean mapFlattened,
byte[] mapKeyNameBytes,
byte[] mapValueNameBytes,
byte[] mapEntryNameBytes,
TimestampFormatTrait.Format timestampFormat) {}

@Override
public SchemaExtensionKey<QueryMemberBinding> key() {
Expand All @@ -42,9 +65,73 @@ public QueryMemberBinding provide(Schema schema) {
return null;
}

byte[] awsName = encodeName(resolveAwsQueryName(schema));
byte[] ec2Name = encodeName(resolveEc2QueryName(schema));

// Pre-compute list metadata if target is a list
boolean listFlattened = false;
byte[] listMemberNameBytes = null;
Schema target = schema.memberTarget();
if (target != null && target.type() == ShapeType.LIST) {
listFlattened = schema.hasTrait(TraitKey.XML_FLATTENED_TRAIT);
if (!listFlattened) {
Schema listMember = target.listMember();
if (listMember != null) {
var xmlName = listMember.getTrait(TraitKey.XML_NAME_TRAIT);
listMemberNameBytes = xmlName != null
? xmlName.getValue().getBytes(StandardCharsets.UTF_8)
: MEMBER_BYTES;
} else {
listMemberNameBytes = MEMBER_BYTES;
}
}
}

// Pre-compute map metadata if target is a map
boolean mapFlattened = false;
byte[] mapKeyNameBytes = null;
byte[] mapValueNameBytes = null;
byte[] mapEntryNameBytes = null;
if (target != null && target.type() == ShapeType.MAP) {
mapFlattened = schema.hasTrait(TraitKey.XML_FLATTENED_TRAIT);
Schema keySchema = target.mapKeyMember();
Schema valueSchema = target.mapValueMember();
if (keySchema != null) {
var keyXmlName = keySchema.getTrait(TraitKey.XML_NAME_TRAIT);
mapKeyNameBytes = keyXmlName != null
? keyXmlName.getValue().getBytes(StandardCharsets.UTF_8)
: KEY_BYTES;
} else {
mapKeyNameBytes = KEY_BYTES;
}
if (valueSchema != null) {
var valueXmlName = valueSchema.getTrait(TraitKey.XML_NAME_TRAIT);
mapValueNameBytes = valueXmlName != null
? valueXmlName.getValue().getBytes(StandardCharsets.UTF_8)
: VALUE_BYTES;
} else {
mapValueNameBytes = VALUE_BYTES;
}
mapEntryNameBytes = mapFlattened ? null : ENTRY_BYTES;
}

// Pre-resolve timestamp format
TimestampFormatTrait.Format timestampFormat = null;
var tsFmt = schema.getTrait(TraitKey.TIMESTAMP_FORMAT_TRAIT);
if (tsFmt != null) {
timestampFormat = tsFmt.getFormat();
}

return new QueryMemberBinding(
encodeName(resolveAwsQueryName(schema)),
encodeName(resolveEc2QueryName(schema)));
awsName,
ec2Name,
listFlattened,
listMemberNameBytes,
mapFlattened,
mapKeyNameBytes,
mapValueNameBytes,
mapEntryNameBytes,
timestampFormat);
}

private static String resolveAwsQueryName(Schema schema) {
Expand All @@ -71,7 +158,7 @@ static byte[] encodeName(String name) {
boolean needsEncoding = false;
for (int i = 0; i < len; i++) {
char c = name.charAt(i);
if (!FormUrlEncodedSink.isUnreserved(c)) {
if (c >= 128 || !QueryFormSerializer.UNRESERVED[c]) {
needsEncoding = true;
break;
}
Expand All @@ -83,11 +170,57 @@ static byte[] encodeName(String name) {
return result;
}

FormUrlEncodedSink tmp = new FormUrlEncodedSink(len * 3);
tmp.writeUrlEncoded(name);
ByteBuffer bb = tmp.finish();
byte[] result = new byte[bb.remaining()];
bb.get(result);
// Member names that need encoding are rare (non-ASCII names).
// Use a simple byte array builder for this cold path.
// Max 12 bytes per char (4-byte UTF-8, each byte percent-encoded to 3 bytes)
byte[] buf = new byte[len * 12];
int pos = 0;
for (int i = 0; i < len; i++) {
char c = name.charAt(i);
if (c < 128 && QueryFormSerializer.UNRESERVED[c]) {
buf[pos++] = (byte) c;
} else if (c < 0x80) {
int off = c * 3;
buf[pos++] = QueryFormSerializer.PERCENT_ENCODED[off];
buf[pos++] = QueryFormSerializer.PERCENT_ENCODED[off + 1];
buf[pos++] = QueryFormSerializer.PERCENT_ENCODED[off + 2];
} else if (c < 0x800) {
int b0 = 0xC0 | (c >> 6);
int b1 = 0x80 | (c & 0x3F);
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b0 * 3, buf, pos, 3);
pos += 3;
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b1 * 3, buf, pos, 3);
pos += 3;
} else if (Character.isHighSurrogate(c) && i + 1 < len
&& Character.isLowSurrogate(name.charAt(i + 1))) {
char low = name.charAt(++i);
int cp = Character.toCodePoint(c, low);
int b0 = 0xF0 | (cp >> 18);
int b1 = 0x80 | ((cp >> 12) & 0x3F);
int b2 = 0x80 | ((cp >> 6) & 0x3F);
int b3 = 0x80 | (cp & 0x3F);
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b0 * 3, buf, pos, 3);
pos += 3;
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b1 * 3, buf, pos, 3);
pos += 3;
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b2 * 3, buf, pos, 3);
pos += 3;
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b3 * 3, buf, pos, 3);
pos += 3;
} else {
int b0 = 0xE0 | (c >> 12);
int b1 = 0x80 | ((c >> 6) & 0x3F);
int b2 = 0x80 | (c & 0x3F);
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b0 * 3, buf, pos, 3);
pos += 3;
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b1 * 3, buf, pos, 3);
pos += 3;
System.arraycopy(QueryFormSerializer.PERCENT_ENCODED, b2 * 3, buf, pos, 3);
pos += 3;
}
}
byte[] result = new byte[pos];
System.arraycopy(buf, 0, result, 0, pos);
return result;
}
}
2 changes: 1 addition & 1 deletion ...src/main/java/software/amazon/smithy/java/aws/client/awsquery/Ec2QueryClientProtocol.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,7 +69,7 @@ public <I extends SerializableStruct, O extends SerializableStruct> HttpRequest
SmithyUri endpoint
) {
String operationName = operation.schema().id().getName();
QueryFormSerializer serializer = new QueryFormSerializer(
QueryFormSerializer serializer = QueryFormSerializer.acquire(
QueryFormSerializer.QueryVariant.EC2_QUERY,
operationName,
version);
Expand Down
Loading