Make TaskUpdate size exceeded message succinct

Improve the exception message for EXCEEDED_TASK_UPDATE_SIZE_LIMIT by
using the DataSize class to produce better human readable succinct values
for byte values.

A message like "TaskUpdate size of 860492511 Bytes has exceeded the limit
of 524288000 Bytes" becomes "TaskUpdate size of 820.63MB has exceeded
the limit of 500MB".

Add unit tests in the TestHttpRemoteTask class to check for the correctly
formatted exception message.

Author: misterjpapa

presto-main/src/main/java/com/facebook/presto/server/remotetask/HttpRemoteTask.java

@@ -75,6 +75,7 @@
 import com.google.common.util.concurrent.ListenableFuture;
 import com.google.common.util.concurrent.SettableFuture;
 import com.sun.management.ThreadMXBean;
+import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import it.unimi.dsi.fastutil.longs.LongArrayList;
 import org.joda.time.DateTime;
@@ -226,6 +227,7 @@ public final class HttpRemoteTask
     private final HandleResolver handleResolver;
     private final int maxTaskUpdateSizeInBytes;
     private final int maxUnacknowledgedSplits;
+    private final DataSize maxTaskUpdateDataSize;
 
     private final TableWriteInfo tableWriteInfo;
 
@@ -325,6 +327,7 @@ public HttpRemoteTask(
             this.handleResolver = handleResolver;
             this.tableWriteInfo = tableWriteInfo;
             this.maxTaskUpdateSizeInBytes = maxTaskUpdateSizeInBytes;
+            this.maxTaskUpdateDataSize = DataSize.succinctBytes(this.maxTaskUpdateSizeInBytes);
             this.maxUnacknowledgedSplits = getMaxUnacknowledgedSplitsPerTask(session);
             checkArgument(maxUnacknowledgedSplits > 0, "maxUnacknowledgedSplits must be > 0, found: %s", maxUnacknowledgedSplits);
 
@@ -882,7 +885,7 @@ private synchronized void sendUpdate()
         taskUpdateRequestSize.add(taskUpdateRequestJson.length);
 
         if (taskUpdateRequestJson.length > maxTaskUpdateSizeInBytes) {
-            failTask(new PrestoException(EXCEEDED_TASK_UPDATE_SIZE_LIMIT, format("TaskUpdate size of %d Bytes has exceeded the limit of %d Bytes", taskUpdateRequestJson.length, maxTaskUpdateSizeInBytes)));
+            failTask(new PrestoException(EXCEEDED_TASK_UPDATE_SIZE_LIMIT, getExceededTaskUpdateSizeMessage(taskUpdateRequestJson)));
         }
 
         if (fragment.isPresent()) {
@@ -928,6 +931,12 @@ private synchronized void sendUpdate()
                 executor);
     }
 
+    private String getExceededTaskUpdateSizeMessage(byte[] taskUpdateRequestJson)
+    {
+        DataSize taskUpdateSize = DataSize.succinctBytes(taskUpdateRequestJson.length);
+        return format("TaskUpdate size of %s has exceeded the limit of %s", taskUpdateSize.toString(), this.maxTaskUpdateDataSize.toString());
+    }
+
     private synchronized List<TaskSource> getSources()
     {
         return Stream.concat(tableScanPlanNodeIds.stream(), remoteSourcePlanNodeIds.stream())

presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTask.java

@@ -72,6 +72,7 @@
 import com.google.inject.Injector;
 import com.google.inject.Module;
 import com.google.inject.Provides;
+import io.airlift.units.DataSize;
 import io.airlift.units.Duration;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -90,6 +91,7 @@
 import javax.ws.rs.core.MediaType;
 import javax.ws.rs.core.UriInfo;
 
+import java.lang.reflect.Method;
 import java.net.URI;
 import java.util.HashMap;
 import java.util.Map;
@@ -222,6 +224,66 @@ public void testHTTPRemoteTaskSize()
         assertTrue(httpRemoteTaskFactory.getTaskUpdateRequestSize() > 0);
     }
 
+    @Test(timeOut = 50000)
+    public void testHTTPRemoteBadTaskSize()
+            throws Exception
+    {
+        AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime());
+        TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, FailureScenario.NO_FAILURE);
+        boolean useThriftEncoding = false;
+        DataSize maxDataSize = DataSize.succinctBytes(1024);
+        InternalCommunicationConfig internalCommunicationConfig = new InternalCommunicationConfig()
+                .setThriftTransportEnabled(useThriftEncoding)
+                .setMaxTaskUpdateSize(maxDataSize);
+
+        HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource, useThriftEncoding, internalCommunicationConfig);
+
+        RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory);
+        testingTaskResource.setInitialTaskInfo(remoteTask.getTaskInfo());
+        remoteTask.start();
+        waitUntilIdle(lastActivityNanos);
+        httpRemoteTaskFactory.stop();
+
+        assertTrue(remoteTask.getTaskStatus().getState().isDone(), format("TaskStatus is not in a done state: %s", remoteTask.getTaskStatus()));
+        assertEquals(getOnlyElement(remoteTask.getTaskStatus().getFailures()).getMessage(), "TaskUpdate size of 1.97kB has exceeded the limit of 1kB");
+    }
+
+    @Test(dataProvider = "getUpdateSize")
+    public void testGetExceededTaskUpdateSizeListMessage(int updateSizeInBytes, int maxDataSizeInBytes,
+                                                         String expectedMessage) throws Exception
+    {
+        AtomicLong lastActivityNanos = new AtomicLong(System.nanoTime());
+        TestingTaskResource testingTaskResource = new TestingTaskResource(lastActivityNanos, FailureScenario.NO_FAILURE);
+        boolean useThriftEncoding = false;
+        DataSize maxDataSize = DataSize.succinctBytes(maxDataSizeInBytes);
+        InternalCommunicationConfig internalCommunicationConfig = new InternalCommunicationConfig()
+                .setThriftTransportEnabled(useThriftEncoding)
+                .setMaxTaskUpdateSize(maxDataSize);
+        HttpRemoteTaskFactory httpRemoteTaskFactory = createHttpRemoteTaskFactory(testingTaskResource, useThriftEncoding, internalCommunicationConfig);
+        RemoteTask remoteTask = createRemoteTask(httpRemoteTaskFactory);
+
+        Method targetMethod = HttpRemoteTask.class.getDeclaredMethod("getExceededTaskUpdateSizeMessage", new Class[]{byte[].class});
+        targetMethod.setAccessible(true);
+        byte[] taskUpdateRequestJson = new byte[updateSizeInBytes];
+        String message = (String) targetMethod.invoke(remoteTask, new Object[]{taskUpdateRequestJson});
+        assertEquals(message, expectedMessage);
+    }
+
+    @DataProvider(name = "getUpdateSize")
+    protected Object[][] getUpdateSize()
+    {
+        return new Object[][] {
+                {2000, 1000, "TaskUpdate size of 1.95kB has exceeded the limit of 1000B"},
+                {2000, 1024, "TaskUpdate size of 1.95kB has exceeded the limit of 1kB"},
+                {5000, 4 * 1024, "TaskUpdate size of 4.88kB has exceeded the limit of 4kB"},
+                {2 * 1024, 1024, "TaskUpdate size of 2kB has exceeded the limit of 1kB"},
+                {1024 * 1024, 512 * 1024, "TaskUpdate size of 1MB has exceeded the limit of 512kB"},
+                {16 * 1024 * 1024, 8 * 1024 * 1024, "TaskUpdate size of 16MB has exceeded the limit of 8MB"},
+                {485 * 1000 * 1000, 1024 * 1024 * 512, "TaskUpdate size of 462.53MB has exceeded the limit of 512MB"},
+                {1024 * 1024 * 1024, 1024 * 1024 * 512, "TaskUpdate size of 1GB has exceeded the limit of 512MB"},
+                {860492511, 524288000, "TaskUpdate size of 820.63MB has exceeded the limit of 500MB"}};
+    }
+
     private void runTest(FailureScenario failureScenario, boolean useThriftEncoding)
             throws Exception
     {
@@ -272,6 +334,13 @@ private RemoteTask createRemoteTask(HttpRemoteTaskFactory httpRemoteTaskFactory)
 
     private static HttpRemoteTaskFactory createHttpRemoteTaskFactory(TestingTaskResource testingTaskResource, boolean useThriftEncoding)
             throws Exception
+    {
+        InternalCommunicationConfig internalCommunicationConfig = new InternalCommunicationConfig().setThriftTransportEnabled(useThriftEncoding);
+        return createHttpRemoteTaskFactory(testingTaskResource, useThriftEncoding, internalCommunicationConfig);
+    }
+
+    private static HttpRemoteTaskFactory createHttpRemoteTaskFactory(TestingTaskResource testingTaskResource, boolean useThriftEncoding, InternalCommunicationConfig internalCommunicationConfig)
+            throws Exception
     {
         Bootstrap app = new Bootstrap(
                 new JsonModule(),
@@ -347,7 +416,7 @@ private HttpRemoteTaskFactory createHttpRemoteTaskFactory(
                                 metadataUpdatesJsonCodec,
                                 metadataUpdatesSmileCodec,
                                 new RemoteTaskStats(),
-                                new InternalCommunicationConfig().setThriftTransportEnabled(useThriftEncoding),
+                                internalCommunicationConfig,
                                 createTestMetadataManager(),
                                 new TestQueryManager(),
                                 new HandleResolver(),