Refactoring for retryability Published on
This story happened when my team received a request that one service was failing to download a PDF file from a remote server. Since this part of the code was very old, there were no metrics/logs or anything to analyse further. It was also unknown when or in which situation the download failed. The only known issue was that the process, including the PDF download, got stuck.
Let's have a look at the existing (old) code which is responsible for downloading the pdf file.
public static String downloadPdf(final String pdfLink, String httpProxy) throws Exception
{
InputStream in = null;
try
{
URLConnection conn = getConnection(pdfLink, httpProxy);
in = conn.getInputStream();
ByteArrayOutputStream tmpOut = new ByteArrayOutputStream();
byte[] buf = new byte[512];
while (true)
{
int len = in.read(buf);
if (len == -1)
{
break;
}
tmpOut.write(buf, 0, len);
}
in.close();
tmpOut.close();
return Base64.getEncoder().encodeToString(tmpOut.toByteArray());
}
catch (Exception e)
{
throw new RuntimeException("error loading pdf");
}
finally
{
if (in != null)
{
try
{
in.close();
}
catch (IOException e)
{
LOG.debug("error closing input stream", e);
}
}
}
}
private URLConnection getConnection(String pdfLink, String httpProxy) throws IOException,
URISyntaxException
{
URLConnection conn;
URL dvUrl = new URL(pdfLink);
if (httpProxy.equals("noProxy"))
{
conn = dvUrl.openConnection();
}
else
{
HttpHost httpHost = HttpHost.create(httpProxy);
InetSocketAddress sa = new InetSocketAddress(httpHost.getHostName(), httpHost.getPort());
conn = dvUrl.openConnection(new Proxy(Proxy.Type.HTTP, sa));
}
return conn;
}
From the first look we can see that this code is very low level, it's creating the connection and the necessary streams to fetch the file.
Even though it is so low-level, this code may look all right, which is largely true, there is just one thing missing: timeouts. Looking at the Javadoc for UrlConnection.setConnectTimeout and UrlConnection.setReadTimeout, we can read the following:
A timeout of zero is interpreted as an infinite timeout.Since we don't set anything in our getConnection method it defaults to 0. Meaning we will wait for infinity and longer!
So the first thing we definitely want to do is set these timeouts.
public static String downloadPdf(final String pdfLink, String httpProxy, int timeout) throws Exception
{
InputStream in = null;
try
{
URLConnection conn = getConnection(pdfLink, httpProxy);
conn.setConnectTimeout(timeout); // !
conn.setReadTimeout(timeout); // !
in = conn.getInputStream();
ByteArrayOutputStream tmpOut = new ByteArrayOutputStream();
...
...One thing I may should have added at the very beginning is a golden rule - whenever there is something not working, an error, a bugfix or an unknown problem: Write a test to confirm the problem. This should be the very first step towards a solution for you. Let the tests be your guide and verify your assumptions.
No test, no fix.
Try to imagine writing a (unit-)test for the downloadPdf method. I couldn't think of a good way, since the whole method relies on a real network connection for execution. This is a big barrier for us, so we need to refactor and create a test for it.
Add Tests
Coming back to our testing rule, we create our first test in such a way that our PDF download is testable without requiring an external connection.
@ExtendWith(MockitoExtension.class)
class PdfDownloadServiceTest {
@Mock
private UrlConnectionPdfClient mockClient;
private PdfDownloadService service;
private String testUrl = "https://example.com/document.pdf";
private String testProxy = "noProxy";
private int timeout = 2000;
@BeforeEach
void setUp() {
service = new PdfDownloadService(mockClient);
}
@Test
void download_successful() throws Exception {
byte[] pdfContent = "Test PDF Content".getBytes();
String expectedBase64 = Base64.getEncoder().encodeToString(pdfContent);
when(mockClient.fetch(testUrl, testProxy, timeout))
.thenReturn(pdfContent);
String result = service.downloadPdf(testUrl, testProxy, timeout);
assertEquals(expectedBase64, result);
verify(mockClient).fetch(testUrl, testProxy, timeout);
}
@Test
void download_fails() throws Exception {
IOException exception = new IOException("Connection timed out");
when(mockClient.fetch(testUrl, testProxy, timeout))
.thenThrow(exception);
CustomException thrown = assertThrows(CustomException.class, () ->
service.downloadPdf(testUrl, testProxy, timeout));
assertTrue(thrown.getMessage().contains("Error loading PDF from"));
assertEquals(exception, thrown.getCause());
verify(mockClient).fetch(testUrl, testProxy, timeout);
}
}As we can see, we already expect to encapsulate the low‑level connection handling through the UrlConnectionPdfClient. Because of this refactoring, we are now able to test our PDF download without any external connection or dependency.
After we are done with the refactoring (extracting the UrlConnectionPdfClient for connection handling) and our test passes, we can take a look at introducing a retry mechanism.
Adding Retries
Not only should we fix the connection issues but we should also add a configurable retry mechanism to this PDF download, since the external server is not maintained by the same company and is obviously unreliable.
As always, start with the tests. We will change our tests to work with a retry config and assert on the retry being done.
public record PdfDownloadRetryConfig(int retryTimeout, int retries, int downloadTimeout) {
}@ExtendWith(MockitoExtension.class)
class PdfDownloadServiceTest {
@Mock
private UrlConnectionPdfClient mockClient;
private PdfDownloadService service;
private PdfDownloadRetryConfig config = new PdfDownloadRetryConfig(1000, 3, 2000);
private String testUrl = "https://example.com/document.pdf";
private String testProxy = "noProxy";
@BeforeEach
void setUp() {
service = new PdfDownloadService(mockClient);
}
@Test
void downloadWithRetry_successful() throws Exception {
byte[] pdfContent = "Test PDF Content".getBytes();
String expectedBase64 = Base64.getEncoder().encodeToString(pdfContent);
when(mockClient.fetch(testUrl, testProxy, config.downloadTimeout()))
.thenReturn(pdfContent);
String result = service.downloadWithRetry(testUrl, testProxy, config);
assertEquals(expectedBase64, result);
verify(mockClient).fetch(testUrl, testProxy, config.downloadTimeout());
}
@Test
void downloadWithRetry_failsAfterRetries() throws Exception {
IOException exception = new IOException("Connection timed out");
when(mockClient.fetch(testUrl, testProxy, config.downloadTimeout()))
.thenThrow(exception);
CustomException thrown = assertThrows(CustomException.class, () ->
service.downloadWithRetry(testUrl, testProxy, config));
assertTrue(thrown.getMessage().contains("Error loading PDF from"));
assertEquals(exception, thrown.getCause());
verify(mockClient, times(config.retries())).fetch(testUrl, testProxy, config.downloadTimeout());
}
}After we are done implementing those changes and turning our tests green, we have arrived at our final implementation, all guided by our tests.
With all the lower level details cut out our new and refactored PdfDownloadService looks like this:
@Service
public class PdfDownloadService
{
private final UrlConnectionPdfClient client;
public PdfDownloadService(UrlConnectionPdfClient client)
{
this.client = client;
}
public String downloadWithRetry(String url, String proxy, PdfDownloadRetryConfig cfg)
throws Exception
{
try
{
return RetryExecutor.builder()
.maxAttempts(cfg.retries())
.backOffPeriod(cfg.retryTimeout())
.build()
.execute(() -> Base64.getEncoder().encodeToString(client.fetch(url, proxy, cfg.downloadTimeout())));
}
catch (Exception e)
{
throw new CustomException(
"Error loading PDF from " + hideCredentialsInUrl(url));
}
}
private String hideCredentialsInUrl(String url)
{
..
..
}
}And the UrlConnectionPdfClient ended up like this:
@Component
public class UrlConnectionPdfClient {
public byte[] fetch(String url, String httpProxy, int timeout) throws IOException, URISyntaxException
{
URLConnection conn = getConnection(url, httpProxy);
conn.setConnectTimeout(timeout);
conn.setReadTimeout(timeout);
try (InputStream in = conn.getInputStream();
ByteArrayOutputStream out = new ByteArrayOutputStream()) {
in.transferTo(out);
return out.toByteArray();
}
}
private URLConnection getConnection(String pdfLink, String httpProxy) throws IOException,
URISyntaxException
{
URLConnection conn;
URL dvUrl = new URL(pdfLink);
if (httpProxy.equals("noProxy"))
{
conn = dvUrl.openConnection();
}
else
{
HttpHost httpHost = HttpHost.create(httpProxy);
InetSocketAddress sa = new InetSocketAddress(httpHost.getHostName(), httpHost.getPort());
conn = dvUrl.openConnection(new Proxy(Proxy.Type.HTTP, sa));
}
return conn;
}
}Conclusion
The main takeaway is to separate the low level details from higher level concerns. This is not always obvious, but it becomes clear when you create good and thorough tests. If there is no obvious or easy way to test something, there is likely a design flaw that should be corrected through refactoring.
For completeness, I am also adding the RetryExecutor (which is basically just a wrapper around RetryTemplate from Spring Retry, you could also just use @Retryable).
public class RetryExecutor {
private static final Logger LOG = LoggerFactory.getLogger(RetryExecutor.class);
private final RetryTemplate retryTemplate;
private RetryExecutor(RetryTemplate retryTemplate) {
this.retryTemplate = retryTemplate;
}
public <T> T execute(Callable<T> operation) throws Exception
{
return retryTemplate.execute(context -> {
LOG.info("Retry attempt: {}", context.getRetryCount());
if (context.getLastThrowable() != null) {
LOG.info("Retrying operation due to: {}", context.getLastThrowable().toString());
}
return operation.call();
});
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private int maxAttempts = 3;
private long backOffPeriod = 1000;
private final Map<Class<? extends Throwable>, Boolean> retryableExceptions = new HashMap<>();
public Builder() {
retryableExceptions.put(Exception.class, true);
}
public Builder maxAttempts(int maxAttempts) {
this.maxAttempts = maxAttempts;
return this;
}
public Builder backOffPeriod(long backOffPeriodMs) {
this.backOffPeriod = backOffPeriodMs;
return this;
}
public Builder retryOn(Class<? extends Throwable> exceptionClass) {
retryableExceptions.clear();
retryableExceptions.put(exceptionClass, true);
return this;
}
public RetryExecutor build() {
SimpleRetryPolicy retryPolicy = new SimpleRetryPolicy(maxAttempts, retryableExceptions);
FixedBackOffPolicy backOffPolicy = new FixedBackOffPolicy();
backOffPolicy.setBackOffPeriod(backOffPeriod);
RetryTemplate template = new RetryTemplate();
template.setRetryPolicy(retryPolicy);
template.setBackOffPolicy(backOffPolicy);
return new RetryExecutor(template);
}
}
}