Skip to content

Commit 45c994f

Browse files
shrinidhijoshishrinidhijoshi
committed
fix: Add createTime and endTime to task RuntimeStats to match native
1 parent a2c5b8f commit 45c994f

File tree

2 files changed

+172
-2
lines changed

2 files changed

+172
-2
lines changed

‎presto-main-base/src/main/java/com/facebook/presto/operator/TaskContext.java‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555

5656
import static com.facebook.airlift.units.DataSize.Unit.BYTE;
5757
import static com.facebook.airlift.units.DataSize.succinctBytes;
58+
import static com.facebook.presto.common.RuntimeUnit.NONE;
5859
import static com.facebook.presto.sql.planner.optimizations.PlanNodeSearcher.searchFrom;
5960
import static com.google.common.base.Preconditions.checkArgument;
6061
import static com.google.common.base.Preconditions.checkState;
@@ -562,12 +563,20 @@ public TaskStats getTaskStats()
562563

563564
boolean fullyBlocked = hasRunningPipelines && runningPipelinesFullyBlocked;
564565

566+
// Add createTime and endTime metrics to RuntimeStats to match native execution behavior
567+
long createTimeInMillis = taskStateMachine.getCreatedTimeInMillis();
568+
long endTimeInMillis = executionEndTime.get();
569+
mergedRuntimeStats.addMetricValue("createTime", NONE, createTimeInMillis);
570+
if (endTimeInMillis > 0) {
571+
mergedRuntimeStats.addMetricValue("endTime", NONE, endTimeInMillis);
572+
}
573+
565574
return new TaskStats(
566-
taskStateMachine.getCreatedTimeInMillis(),
575+
createTimeInMillis,
567576
executionStartTime.get(),
568577
lastExecutionStartTime.get(),
569578
lastExecutionEndTime,
570-
executionEndTime.get(),
579+
endTimeInMillis,
571580
elapsedTimeInNanos,
572581
queuedTimeInNanos,
573582
totalDrivers,

‎presto-main-base/src/test/java/com/facebook/presto/operator/TestTaskContextRuntimeStats.java‎

Lines changed: 161 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,161 @@
1+
/*
2+
* Licensed under the Apache License, Version 2.0 (the "License");
3+
* you may not use this file except in compliance with the License.
4+
* You may obtain a copy of the License at
5+
*
6+
* http://www.apache.org/licenses/LICENSE-2.0
7+
*
8+
* Unless required by applicable law or agreed to in writing, software
9+
* distributed under the License is distributed on an "AS IS" BASIS,
10+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
* See the License for the specific language governing permissions and
12+
* limitations under the License.
13+
*/
14+
package com.facebook.presto.operator;
15+
16+
import com.facebook.airlift.stats.TestingGcMonitor;
17+
import com.facebook.presto.Session;
18+
import com.facebook.presto.common.RuntimeStats;
19+
import com.facebook.presto.execution.StageExecutionId;
20+
import com.facebook.presto.execution.StageId;
21+
import com.facebook.presto.execution.TaskId;
22+
import com.facebook.presto.execution.TaskStateMachine;
23+
import com.facebook.presto.memory.QueryContext;
24+
import com.facebook.presto.spi.QueryId;
25+
import com.facebook.presto.spi.memory.MemoryPoolId;
26+
import com.facebook.presto.spiller.SpillSpaceTracker;
27+
import com.google.common.util.concurrent.MoreExecutors;
28+
import org.testng.annotations.AfterClass;
29+
import org.testng.annotations.Test;
30+
31+
import java.util.Optional;
32+
import java.util.concurrent.ScheduledExecutorService;
33+
34+
import static com.facebook.airlift.concurrent.Threads.daemonThreadsNamed;
35+
import static com.facebook.airlift.json.JsonCodec.listJsonCodec;
36+
import static com.facebook.airlift.units.DataSize.Unit.GIGABYTE;
37+
import static com.facebook.airlift.units.DataSize.Unit.MEGABYTE;
38+
import static com.facebook.airlift.units.DataSize.succinctBytes;
39+
import static com.facebook.presto.testing.TestingSession.testSessionBuilder;
40+
import static java.util.concurrent.Executors.newScheduledThreadPool;
41+
import static org.testng.Assert.assertEquals;
42+
import static org.testng.Assert.assertNotNull;
43+
import static org.testng.Assert.assertTrue;
44+
45+
public class TestTaskContextRuntimeStats
46+
{
47+
private final ScheduledExecutorService scheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("test-%s"));
48+
49+
@AfterClass(alwaysRun = true)
50+
public void tearDown()
51+
{
52+
scheduledExecutor.shutdownNow();
53+
}
54+
55+
@Test
56+
public void testTaskStatsIncludesCreateTimeAndEndTime()
57+
{
58+
Session session = testSessionBuilder().build();
59+
QueryContext queryContext = createQueryContext(session);
60+
61+
TaskStateMachine taskStateMachine = new TaskStateMachine(
62+
new TaskId(new StageExecutionId(new StageId(new QueryId("test_query"), 0), 0), 0, 0),
63+
MoreExecutors.directExecutor());
64+
65+
TaskContext taskContext = queryContext.addTaskContext(
66+
taskStateMachine,
67+
session,
68+
Optional.empty(),
69+
false,
70+
false,
71+
false,
72+
false,
73+
false);
74+
75+
long createTimeBeforeStats = taskStateMachine.getCreatedTimeInMillis();
76+
77+
// Get task stats
78+
TaskStats taskStats = taskContext.getTaskStats();
79+
80+
// Verify RuntimeStats contains createTime
81+
RuntimeStats runtimeStats = taskStats.getRuntimeStats();
82+
assertNotNull(runtimeStats, "RuntimeStats should not be null");
83+
assertTrue(runtimeStats.getMetrics().containsKey("createTime"), "RuntimeStats should contain createTime metric");
84+
85+
// Verify createTime value is reasonable
86+
long createTimeFromStats = (long) runtimeStats.getMetric("createTime").getSum();
87+
assertEquals(createTimeFromStats, createTimeBeforeStats, "createTime should match task creation time");
88+
89+
// Mark task as finished to trigger endTime
90+
taskStateMachine.finished();
91+
TaskStats finalTaskStats = taskContext.getTaskStats();
92+
RuntimeStats finalRuntimeStats = finalTaskStats.getRuntimeStats();
93+
94+
// Verify endTime is now present
95+
assertTrue(finalRuntimeStats.getMetrics().containsKey("endTime"), "RuntimeStats should contain endTime metric after task finishes");
96+
long endTimeFromStats = (long) finalRuntimeStats.getMetric("endTime").getSum();
97+
assertTrue(endTimeFromStats > 0, "endTime should be greater than 0");
98+
assertTrue(endTimeFromStats >= createTimeFromStats, "endTime should be >= createTime");
99+
}
100+
101+
@Test
102+
public void testTaskStatsRuntimeStatsNotNullBeforeTaskFinish()
103+
{
104+
Session session = testSessionBuilder().build();
105+
QueryContext queryContext = createQueryContext(session);
106+
107+
TaskStateMachine taskStateMachine = new TaskStateMachine(
108+
new TaskId(new StageExecutionId(new StageId(new QueryId("test_query_2"), 0), 0), 0, 0),
109+
MoreExecutors.directExecutor());
110+
111+
TaskContext taskContext = queryContext.addTaskContext(
112+
taskStateMachine,
113+
session,
114+
Optional.empty(),
115+
false,
116+
false,
117+
false,
118+
false,
119+
false);
120+
121+
// Get stats before task finishes
122+
TaskStats taskStats = taskContext.getTaskStats();
123+
RuntimeStats runtimeStats = taskStats.getRuntimeStats();
124+
125+
// Verify RuntimeStats is not null and contains createTime even before task finishes
126+
assertNotNull(runtimeStats, "RuntimeStats should not be null");
127+
assertTrue(runtimeStats.getMetrics().containsKey("createTime"), "RuntimeStats should contain createTime even before task finishes");
128+
129+
// endTime should not be present yet (or be 0)
130+
if (runtimeStats.getMetrics().containsKey("endTime")) {
131+
long endTime = (long) runtimeStats.getMetric("endTime").getSum();
132+
assertEquals(endTime, 0L, "endTime should be 0 before task finishes");
133+
}
134+
}
135+
136+
private QueryContext createQueryContext(Session session)
137+
{
138+
return new QueryContext(
139+
session.getQueryId(),
140+
succinctBytes(1, MEGABYTE),
141+
succinctBytes(1, GIGABYTE),
142+
succinctBytes(1, MEGABYTE),
143+
succinctBytes(1, GIGABYTE),
144+
new TestingMemoryPool(succinctBytes(1, GIGABYTE)),
145+
new TestingGcMonitor(),
146+
MoreExecutors.directExecutor(),
147+
scheduledExecutor,
148+
succinctBytes(1, GIGABYTE),
149+
new SpillSpaceTracker(succinctBytes(1, GIGABYTE)),
150+
listJsonCodec(TaskMemoryReservationSummary.class));
151+
}
152+
153+
private static class TestingMemoryPool
154+
extends com.facebook.presto.memory.MemoryPool
155+
{
156+
public TestingMemoryPool(com.facebook.airlift.units.DataSize maxMemory)
157+
{
158+
super(new MemoryPoolId("test"), maxMemory);
159+
}
160+
}
161+
}

0 commit comments

Comments
 (0)