Skip to content

Commit 8860e93

Browse files
committed
Add CompressingAsyncEntityProducer for transparent request-side compression via ContentCodecRegistry
1 parent 179491f commit 8860e93

3 files changed

Lines changed: 408 additions & 0 deletions

File tree

Lines changed: 181 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,181 @@
1+
/*
2+
* ====================================================================
3+
* Licensed to the Apache Software Foundation (ASF) under one
4+
* or more contributor license agreements. See the NOTICE file
5+
* distributed with this work for additional information
6+
* regarding copyright ownership. The ASF licenses this file
7+
* to you under the Apache License, Version 2.0 (the
8+
* "License"); you may not use this file except in compliance
9+
* with the License. You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing,
14+
* software distributed under the License is distributed on an
15+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16+
* KIND, either express or implied. See the License for the
17+
* specific language governing permissions and limitations
18+
* under the License.
19+
* ====================================================================
20+
*
21+
* This software consists of voluntary contributions made by many
22+
* individuals on behalf of the Apache Software Foundation. For more
23+
* information on the Apache Software Foundation, please see
24+
* <http://www.apache.org/>.
25+
*
26+
*/
27+
package org.apache.hc.client5.http.async.methods;
28+
29+
import java.io.ByteArrayOutputStream;
30+
import java.io.IOException;
31+
import java.nio.ByteBuffer;
32+
import java.util.Collections;
33+
import java.util.HashSet;
34+
import java.util.List;
35+
import java.util.Locale;
36+
import java.util.Set;
37+
import java.util.concurrent.atomic.AtomicInteger;
38+
39+
import org.apache.hc.client5.http.entity.compress.ContentCodecRegistry;
40+
import org.apache.hc.client5.http.entity.compress.ContentCoding;
41+
import org.apache.hc.core5.http.ContentType;
42+
import org.apache.hc.core5.http.Header;
43+
import org.apache.hc.core5.http.HttpEntity;
44+
import org.apache.hc.core5.http.io.entity.ByteArrayEntity;
45+
import org.apache.hc.core5.http.nio.AsyncEntityProducer;
46+
import org.apache.hc.core5.http.nio.DataStreamChannel;
47+
import org.apache.hc.core5.util.Args;
48+
49+
/**
50+
* Generic {@link AsyncEntityProducer} that compresses the output produced by a
51+
* delegate using any codec available via {@link ContentCodecRegistry#encoder}.
52+
*
53+
* <p>The delegate’s entire payload is buffered once; therefore use this class
54+
* only for modest payload sizes.</p>
55+
*
56+
* @since 5.6
57+
*/
58+
public final class CompressingAsyncEntityProducer implements AsyncEntityProducer {
59+
60+
private final byte[] compressed;
61+
private final String codingToken;
62+
private final ContentType contentType;
63+
private final AtomicInteger cursor = new AtomicInteger();
64+
65+
public CompressingAsyncEntityProducer(
66+
final AsyncEntityProducer delegate,
67+
final String coding) throws IOException {
68+
69+
Args.notNull(delegate, "delegate");
70+
this.codingToken = Args.notBlank(coding, "coding").toLowerCase(Locale.ROOT);
71+
this.contentType = ContentType.parse(delegate.getContentType());
72+
73+
final ByteArrayOutputStream rawBuf = new ByteArrayOutputStream();
74+
delegate.produce(new BufferingChannel(rawBuf));
75+
76+
final HttpEntity rawEntity = new ByteArrayEntity(rawBuf.toByteArray(), contentType);
77+
final HttpEntity encodedEnt = encode(rawEntity, codingToken);
78+
79+
final ByteArrayOutputStream encBuf = new ByteArrayOutputStream();
80+
encodedEnt.writeTo(encBuf);
81+
this.compressed = encBuf.toByteArray();
82+
}
83+
84+
@Override
85+
public boolean isRepeatable() {
86+
return false;
87+
}
88+
89+
@Override
90+
public long getContentLength() {
91+
return compressed.length;
92+
}
93+
94+
@Override
95+
public String getContentType() {
96+
return contentType.toString();
97+
}
98+
99+
@Override
100+
public String getContentEncoding() {
101+
return codingToken;
102+
}
103+
104+
@Override
105+
public boolean isChunked() {
106+
return false;
107+
}
108+
109+
@Override
110+
public int available() {
111+
return compressed.length - cursor.get();
112+
}
113+
114+
@Override
115+
public Set<String> getTrailerNames() {
116+
return new HashSet<>();
117+
}
118+
119+
@Override
120+
public void failed(final Exception ex) {
121+
}
122+
123+
@Override
124+
public void releaseResources() {
125+
}
126+
127+
@Override
128+
public void produce(final DataStreamChannel ch) throws IOException {
129+
final int pos = cursor.get();
130+
if (pos >= compressed.length) {
131+
ch.endStream(Collections.<Header>emptyList());
132+
return;
133+
}
134+
final int chunk = Math.min(8 * 1024, compressed.length - pos);
135+
ch.write(ByteBuffer.wrap(compressed, pos, chunk));
136+
cursor.addAndGet(chunk);
137+
}
138+
139+
private static HttpEntity encode(final HttpEntity src, final String token) throws IOException {
140+
final ContentCoding coding = ContentCoding.fromToken(token);
141+
if (coding == null) {
142+
throw new IOException("Unknown coding: " + token);
143+
}
144+
final java.util.function.UnaryOperator<HttpEntity> op = ContentCodecRegistry.encoder(coding);
145+
if (op == null) {
146+
throw new IOException("No encoder registered for " + token);
147+
}
148+
return op.apply(src);
149+
}
150+
151+
/**
152+
* Captures bytes pushed by the delegate into a buffer.
153+
*/
154+
private static final class BufferingChannel implements DataStreamChannel {
155+
private final ByteArrayOutputStream buf;
156+
157+
BufferingChannel(final ByteArrayOutputStream buf) {
158+
this.buf = buf;
159+
}
160+
161+
@Override
162+
public void requestOutput() {
163+
}
164+
165+
@Override
166+
public int write(final ByteBuffer src) {
167+
final byte[] tmp = new byte[src.remaining()];
168+
src.get(tmp);
169+
buf.write(tmp, 0, tmp.length);
170+
return tmp.length;
171+
}
172+
173+
@Override
174+
public void endStream() {
175+
}
176+
177+
@Override
178+
public void endStream(final List<? extends Header> t) {
179+
}
180+
}
181+
}
Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
/*
2+
* ====================================================================
3+
* Licensed to the Apache Software Foundation (ASF) under one
4+
* or more contributor license agreements. See the NOTICE file
5+
* distributed with this work for additional information
6+
* regarding copyright ownership. The ASF licenses this file
7+
* to you under the Apache License, Version 2.0 (the
8+
* "License"); you may not use this file except in compliance
9+
* with the License. You may obtain a copy of the License at
10+
*
11+
* http://www.apache.org/licenses/LICENSE-2.0
12+
*
13+
* Unless required by applicable law or agreed to in writing,
14+
* software distributed under the License is distributed on an
15+
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
16+
* KIND, either express or implied. See the License for the
17+
* specific language governing permissions and limitations
18+
* under the License.
19+
* ====================================================================
20+
*
21+
* This software consists of voluntary contributions made by many
22+
* individuals on behalf of the Apache Software Foundation. For more
23+
* information on the Apache Software Foundation, please see
24+
* <http://www.apache.org/>.
25+
*
26+
*/
27+
package org.apache.hc.client5.http.async.methods;
28+
29+
import static org.junit.jupiter.api.Assertions.assertEquals;
30+
31+
import java.io.ByteArrayInputStream;
32+
import java.io.ByteArrayOutputStream;
33+
import java.io.IOException;
34+
import java.nio.ByteBuffer;
35+
import java.nio.charset.StandardCharsets;
36+
import java.util.List;
37+
import java.util.zip.GZIPInputStream;
38+
39+
import org.apache.hc.core5.http.ContentType;
40+
import org.apache.hc.core5.http.Header;
41+
import org.apache.hc.core5.http.nio.AsyncEntityProducer;
42+
import org.apache.hc.core5.http.nio.DataStreamChannel;
43+
import org.apache.hc.core5.http.nio.entity.StringAsyncEntityProducer;
44+
import org.junit.jupiter.api.Test;
45+
46+
class TestCompressingAsyncEntityProducer {
47+
48+
49+
@Test
50+
void gzipRoundTrip() throws Exception {
51+
final String plain = "lorem ipsum äëïöü";
52+
53+
final AsyncEntityProducer original =
54+
new StringAsyncEntityProducer(plain, ContentType.TEXT_PLAIN);
55+
56+
final CompressingAsyncEntityProducer gzip =
57+
new CompressingAsyncEntityProducer(original, "gzip");
58+
59+
final ByteArrayOutputStream wire = new ByteArrayOutputStream();
60+
final DataStreamChannel channel = new BufferingChannel(wire);
61+
62+
while (gzip.available() > 0) {
63+
gzip.produce(channel);
64+
}
65+
66+
final String roundTrip = inflateUtf8(wire.toByteArray());
67+
assertEquals(plain, roundTrip);
68+
69+
assertEquals("gzip", gzip.getContentEncoding()); // meta-check
70+
final byte[] wireBytes = wire.toByteArray(); // <-- here
71+
assertEquals(0x1F, wireBytes[0] & 0xFF); // magic bytes
72+
assertEquals(0x8B, wireBytes[1] & 0xFF);
73+
74+
75+
}
76+
77+
private static String inflateUtf8(final byte[] gz) throws IOException {
78+
try (final GZIPInputStream gis = new GZIPInputStream(new ByteArrayInputStream(gz));
79+
final ByteArrayOutputStream out = new ByteArrayOutputStream()) {
80+
81+
final byte[] buf = new byte[8 * 1024];
82+
int n;
83+
while ((n = gis.read(buf)) != -1) {
84+
out.write(buf, 0, n);
85+
}
86+
return out.toString(StandardCharsets.UTF_8.name());
87+
}
88+
}
89+
90+
private static final class BufferingChannel implements DataStreamChannel {
91+
private final ByteArrayOutputStream buf;
92+
93+
BufferingChannel(final ByteArrayOutputStream buf) {
94+
this.buf = buf;
95+
}
96+
97+
@Override
98+
public void requestOutput() {
99+
}
100+
101+
@Override
102+
public int write(final ByteBuffer src) {
103+
final byte[] b = new byte[src.remaining()];
104+
src.get(b);
105+
buf.write(b, 0, b.length);
106+
return b.length;
107+
}
108+
109+
@Override
110+
public void endStream() {
111+
}
112+
113+
@Override
114+
public void endStream(final List<? extends Header> t) {
115+
}
116+
}
117+
}

0 commit comments

Comments
 (0)