Skip to content

Commit 3b47c49

Browse files
bryanckdanielcweeks
authored andcommitted
AWS: abort S3 input stream on close if not EOS (#7262)
* AWS: abort S3 input stream on close if not EOS * Close the stream for backwards compatibility * undo unrelated change * add trace log * comment update * logger updates * handle connection closed exception
1 parent 29e60c4 commit 3b47c49

File tree

1 file changed

+24
-1
lines changed

1 file changed

+24
-1
lines changed

aws/src/main/java/org/apache/iceberg/aws/s3/S3InputStream.java

Lines changed: 24 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.slf4j.Logger;
3737
import org.slf4j.LoggerFactory;
3838
import software.amazon.awssdk.core.sync.ResponseTransformer;
39+
import software.amazon.awssdk.http.Abortable;
3940
import software.amazon.awssdk.services.s3.S3Client;
4041
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
4142
import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
@@ -196,7 +197,29 @@ private void openStream() throws IOException {
196197

197198
private void closeStream() throws IOException {
198199
if (stream != null) {
199-
stream.close();
200+
// if we aren't at the end of the stream, and the stream is abortable, then
201+
// call abort() so we don't read the remaining data with the Apache HTTP client
202+
abortStream();
203+
try {
204+
stream.close();
205+
} catch (IOException e) {
206+
// the Apache HTTP client will throw a ConnectionClosedException
207+
// when closing an aborted stream, which is expected
208+
if (!e.getClass().getSimpleName().equals("ConnectionClosedException")) {
209+
throw e;
210+
}
211+
}
212+
stream = null;
213+
}
214+
}
215+
216+
private void abortStream() {
217+
try {
218+
if (stream instanceof Abortable && stream.read() != -1) {
219+
((Abortable) stream).abort();
220+
}
221+
} catch (Exception e) {
222+
LOG.warn("An error occurred while aborting the stream", e);
200223
}
201224
}
202225

0 commit comments

Comments
 (0)