RestTemplate下载大文件时OOM问题解决

1 背景

代码中使用RestTemplate下载大文件，发现会OOM，代码如下：

RestTemplate restTemplate = new RestTemplate();

// 会OOM
ResponseEntity<byte[]> entity = restTemplate.getForEntity("http://localhost:8088/1.jar", byte[].class);
log.info(entity.toString());

报错信息如下：

Exception in thread "main" java.lang.OutOfMemoryError: Java heap space
	at java.io.ByteArrayOutputStream.<init>(ByteArrayOutputStream.java:77)
	at org.springframework.http.converter.ByteArrayHttpMessageConverter.readInternal(ByteArrayHttpMessageConverter.java:57)
	at org.springframework.http.converter.ByteArrayHttpMessageConverter.readInternal(ByteArrayHttpMessageConverter.java:39)
	at org.springframework.http.converter.AbstractHttpMessageConverter.read(AbstractHttpMessageConverter.java:199)
	at org.springframework.web.client.HttpMessageConverterExtractor.extractData(HttpMessageConverterExtractor.java:114)
	at org.springframework.web.client.RestTemplate$ResponseEntityResponseExtractor.extractData(RestTemplate.java:996)
	at org.springframework.web.client.RestTemplate$ResponseEntityResponseExtractor.extractData(RestTemplate.java:979)
	at org.springframework.web.client.RestTemplate.doExecute(RestTemplate.java:739)
	at org.springframework.web.client.RestTemplate.execute(RestTemplate.java:672)
	at org.springframework.web.client.RestTemplate.getForEntity(RestTemplate.java:340)
	at com.xxx.demo.DemoApplication.main(DemoApplication.java:40)

2 解决

RestTemplate restTemplate = new RestTemplate();

// 直接写入文件，不会OOM
RequestCallback requestCallback = restTemplate.acceptHeaderRequestCallback(byte[].class);
restTemplate.execute("http://localhost:8088/1.jar", HttpMethod.GET, requestCallback, response -> {
    // 示例，写入到D盘
    IOUtils.copyLarge(response.getBody(), FileUtils.openOutputStream(new File("D:\\1.jar")));
    return null;
});

3 源码分析

让我们来看看RestTemplate的getForEntity为什么会OOM？内存不够，无法分配，导致OOM。

因为要返回的类型为byts[]，所以RestTemplate使用ByteArrayHttpMessageConverter来读取Response，并将Response转换为byte[]。该类源码如下：

public class ByteArrayHttpMessageConverter extends AbstractHttpMessageConverter<byte[]> {

	public ByteArrayHttpMessageConverter() {
		super(MediaType.APPLICATION_OCTET_STREAM, MediaType.ALL);
	}

        // 表示支持对byte[]进行转换
	@Override
	public boolean supports(Class<?> clazz) {
		return byte[].class == clazz;
	}

	@Override
	public byte[] readInternal(Class<? extends byte[]> clazz, HttpInputMessage inputMessage) throws IOException {
                // 读取返回的内容长度
		long contentLength = inputMessage.getHeaders().getContentLength();
                // 创建ByteArrayOutputStream（关键报错点！）
		ByteArrayOutputStream bos =
				new ByteArrayOutputStream(contentLength >= 0 ? (int) contentLength : StreamUtils.BUFFER_SIZE);
                // 将response中内容复制到字节数组输出流中
		StreamUtils.copy(inputMessage.getBody(), bos);
                // 将字节数组输出流转换为字节数组，并返回
		return bos.toByteArray();
	}

	@Override
	protected Long getContentLength(byte[] bytes, @Nullable MediaType contentType) {
		return (long) bytes.length;
	}

	@Override
	protected void writeInternal(byte[] bytes, HttpOutputMessage outputMessage) throws IOException {
		StreamUtils.copy(bytes, outputMessage.getBody());
	}

}

关键的OOM报错点为：

// 创建ByteArrayOutputStream（关键报错点！）
ByteArrayOutputStream bos =
    new ByteArrayOutputStream(contentLength >= 0 ? (int) contentLength : StreamUtils.BUFFER_SIZE);

// 可以看到，这个构造函数，会分配一个size大小的byte[]
// 所以如果此时的JVM堆内存不足以分配，就会抛出OOM异常
public ByteArrayOutputStream(int size) {
    if (size < 0) {
        throw new IllegalArgumentException("Negative initial size: "
                                           + size);
    }
    buf = new byte[size];
}

参考

1. [RestTemplate大文件下载](https://www.cnblogs.com/zimug/archive/2020/08/12/13488517.html)