Add support to disable task update size tracking

Author: NikhilCollooru

presto-main/src/main/java/com/facebook/presto/execution/TaskManagerConfig.java

@@ -51,6 +51,7 @@ public class TaskManagerConfig
     private boolean statisticsCpuTimerEnabled = true;
     private boolean perOperatorAllocationTrackingEnabled;
     private boolean taskAllocationTrackingEnabled;
+    private boolean taskUpdateSizeTrackingEnabled = true;
     private DataSize maxPartialAggregationMemoryUsage = new DataSize(16, Unit.MEGABYTE);
     private DataSize maxLocalExchangeBufferSize = new DataSize(32, Unit.MEGABYTE);
     private DataSize maxIndexMemoryUsage = new DataSize(64, Unit.MEGABYTE);
@@ -666,4 +667,16 @@ public TaskManagerConfig setHighMemoryTaskKillerStrategy(HighMemoryTaskKillerStr
         this.highMemoryTaskKillerStrategy = highMemoryTaskKillerStrategy;
         return this;
     }
+
+    public boolean isTaskUpdateSizeTrackingEnabled()
+    {
+        return taskUpdateSizeTrackingEnabled;
+    }
+
+    @Config("task.update-size-tracking-enabled")
+    public TaskManagerConfig setTaskUpdateSizeTrackingEnabled(boolean taskUpdateSizeTrackingEnabled)
+    {
+        this.taskUpdateSizeTrackingEnabled = taskUpdateSizeTrackingEnabled;
+        return this;
+    }
 }

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

@@ -227,6 +227,7 @@ public final class HttpRemoteTask
     private final TableWriteInfo tableWriteInfo;
 
     private final DecayCounter taskUpdateRequestSize;
+    private final boolean taskUpdateSizeTrackingEnabled;
     private final SchedulerStatsTracker schedulerStatsTracker;
 
     private final HttpRemoteTaskFactory.SafeEventLoop taskEventLoop;
@@ -263,6 +264,7 @@ public static HttpRemoteTask createHttpRemoteTask(
             MetadataManager metadataManager,
             QueryManager queryManager,
             DecayCounter taskUpdateRequestSize,
+            boolean taskUpdateSizeTrackingEnabled,
             HandleResolver handleResolver,
             ConnectorTypeSerdeManager connectorTypeSerdeManager,
             SchedulerStatsTracker schedulerStatsTracker,
@@ -299,6 +301,7 @@ public static HttpRemoteTask createHttpRemoteTask(
                 metadataManager,
                 queryManager,
                 taskUpdateRequestSize,
+                taskUpdateSizeTrackingEnabled,
                 handleResolver,
                 connectorTypeSerdeManager,
                 schedulerStatsTracker,
@@ -338,6 +341,7 @@ private HttpRemoteTask(Session session,
             MetadataManager metadataManager,
             QueryManager queryManager,
             DecayCounter taskUpdateRequestSize,
+            boolean taskUpdateSizeTrackingEnabled,
             HandleResolver handleResolver,
             ConnectorTypeSerdeManager connectorTypeSerdeManager,
             SchedulerStatsTracker schedulerStatsTracker,
@@ -404,6 +408,7 @@ private HttpRemoteTask(Session session,
                 .map(PlanNode::getId)
                 .collect(toImmutableSet());
         this.taskUpdateRequestSize = taskUpdateRequestSize;
+        this.taskUpdateSizeTrackingEnabled = taskUpdateSizeTrackingEnabled;
         this.schedulerStatsTracker = schedulerStatsTracker;
 
         for (Entry<PlanNodeId, Split> entry : requireNonNull(initialSplits, "initialSplits is null").entries()) {
@@ -983,20 +988,22 @@ private void sendUpdate()
             byte[] taskUpdateRequestJson = taskUpdateRequestCodec.toBytes(updateRequest);
             schedulerStatsTracker.recordTaskUpdateSerializedCpuTime(THREAD_MX_BEAN.getCurrentThreadCpuTime() - serializeStartCpuTimeNanos);
 
-            taskUpdateRequestSize.add(taskUpdateRequestJson.length);
-
             if (taskUpdateRequestJson.length > maxTaskUpdateSizeInBytes) {
                 failTask(new PrestoException(EXCEEDED_TASK_UPDATE_SIZE_LIMIT, getExceededTaskUpdateSizeMessage(taskUpdateRequestJson)));
                 return;
             }
 
-            if (fragment.isPresent()) {
-                stats.updateWithPlanSize(taskUpdateRequestJson.length);
-            }
-            else {
-                if (ThreadLocalRandom.current().nextDouble() < UPDATE_WITHOUT_PLAN_STATS_SAMPLE_RATE) {
-                    // This is to keep track of the task update size even when the plan fragment is NOT present
-                    stats.updateWithoutPlanSize(taskUpdateRequestJson.length);
+            if (taskUpdateSizeTrackingEnabled) {
+                taskUpdateRequestSize.add(taskUpdateRequestJson.length);
+
+                if (fragment.isPresent()) {
+                    stats.updateWithPlanSize(taskUpdateRequestJson.length);
+                }
+                else {
+                    if (ThreadLocalRandom.current().nextDouble() < UPDATE_WITHOUT_PLAN_STATS_SAMPLE_RATE) {
+                        // This is to keep track of the task update size even when the plan fragment is NOT present
+                        stats.updateWithoutPlanSize(taskUpdateRequestJson.length);
+                    }
                 }
             }
 

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

@@ -94,6 +94,7 @@ public class HttpRemoteTaskFactory
     private final MetadataManager metadataManager;
     private final QueryManager queryManager;
     private final DecayCounter taskUpdateRequestSize;
+    private final boolean taskUpdateSizeTrackingEnabled;
     private final EventLoopGroup eventLoopGroup = new DefaultEventLoopGroup(Runtime.getRuntime().availableProcessors(),
             new ThreadFactoryBuilder().setNameFormat("task-event-loop-%s").setDaemon(true).build())
     {
@@ -180,6 +181,7 @@ else if (binaryTransportEnabled) {
         this.queryManager = queryManager;
 
         this.taskUpdateRequestSize = new DecayCounter(ExponentialDecay.oneMinute());
+        this.taskUpdateSizeTrackingEnabled = taskConfig.isTaskUpdateSizeTrackingEnabled();
     }
 
     @Managed
@@ -239,6 +241,7 @@ public RemoteTask createRemoteTask(
                 metadataManager,
                 queryManager,
                 taskUpdateRequestSize,
+                taskUpdateSizeTrackingEnabled,
                 handleResolver,
                 connectorTypeSerdeManager,
                 schedulerStatsTracker,

presto-main/src/test/java/com/facebook/presto/execution/TestTaskManagerConfig.java

@@ -77,7 +77,8 @@ public void testDefaults()
                 .setHighMemoryTaskKillerStrategy(HighMemoryTaskKillerStrategy.FREE_MEMORY_ON_FULL_GC)
                 .setHighMemoryTaskKillerGCReclaimMemoryThreshold(0.01)
                 .setHighMemoryTaskKillerFrequentFullGCDurationThreshold(new Duration(1, SECONDS))
-                .setHighMemoryTaskKillerHeapMemoryThreshold(0.9));
+                .setHighMemoryTaskKillerHeapMemoryThreshold(0.9)
+                .setTaskUpdateSizeTrackingEnabled(true));
     }
 
     @Test
@@ -125,6 +126,7 @@ public void testExplicitPropertyMappings()
                 .put("experimental.task.high-memory-task-killer-reclaim-memory-threshold", "0.8")
                 .put("experimental.task.high-memory-task-killer-frequent-full-gc-duration-threshold", "2s")
                 .put("experimental.task.high-memory-task-killer-heap-memory-threshold", "0.8")
+                .put("task.update-size-tracking-enabled", "false")
                 .build();
 
         TaskManagerConfig expected = new TaskManagerConfig()
@@ -168,7 +170,8 @@ public void testExplicitPropertyMappings()
                 .setHighMemoryTaskKillerStrategy(HighMemoryTaskKillerStrategy.FREE_MEMORY_ON_FREQUENT_FULL_GC)
                 .setHighMemoryTaskKillerGCReclaimMemoryThreshold(0.8)
                 .setHighMemoryTaskKillerFrequentFullGCDurationThreshold(new Duration(2, SECONDS))
-                .setHighMemoryTaskKillerHeapMemoryThreshold(0.8);
+                .setHighMemoryTaskKillerHeapMemoryThreshold(0.8)
+                .setTaskUpdateSizeTrackingEnabled(false);
 
         assertFullMapping(properties, expected);
     }