Skip to content

Commit a7a0098

Browse files
agrawal-siddharthagrawal-siddharth
authored
feat: populate client_stats in AppendRowsRequest (#14614)
1 parent 44a12ef commit a7a0098

2 files changed

Lines changed: 196 additions & 0 deletions

File tree

‎java-bigquerystorage/google-cloud-bigquerystorage/src/main/java/com/google/cloud/bigquery/storage/v1/ConnectionWorker.java‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1169,6 +1169,8 @@ private void appendLoop() {
11691169
boolean firstRequestForTableOrSchemaSwitch = true;
11701170

11711171
while (!waitingQueueDrained()) {
1172+
long queuedRequestsCount = 0;
1173+
ClientStats.WindowStats windowStats = null;
11721174
this.lock.lock();
11731175
try {
11741176
hasMessageInWaitingQueue.await(100, TimeUnit.MILLISECONDS);
@@ -1216,6 +1218,11 @@ private void appendLoop() {
12161218
localQueue.addLast(requestWrapper);
12171219
healthCheckMetrics.updateRequestsSent(requestWrapper.messageSize);
12181220
}
1221+
if (!localQueue.isEmpty()) {
1222+
queuedRequestsCount =
1223+
inflightRequestQueue.size(); // waitingRequestQueue is empty so can be omitted here
1224+
windowStats = createWindowStats();
1225+
}
12191226
} catch (InterruptedException e) {
12201227
log.warning(
12211228
"Interrupted while waiting for message. Stream: "
@@ -1302,6 +1309,13 @@ private void appendLoop() {
13021309
}
13031310
firstRequestForTableOrSchemaSwitch = false;
13041311

1312+
ClientStats clientStats =
1313+
createClientStats(wrapper.requestSendTimeStamp, queuedRequestsCount, windowStats);
1314+
windowStats = null;
1315+
if (!clientStats.equals(ClientStats.getDefaultInstance())) {
1316+
originalRequestBuilder.setClientStats(clientStats);
1317+
}
1318+
13051319
requestProfilerHook.startOperation(
13061320
RequestProfiler.OperationName.RESPONSE_LATENCY, requestUniqueId);
13071321

@@ -1320,6 +1334,69 @@ private void appendLoop() {
13201334
cleanupConnectionAndRequests(/* avoidBlocking= */ false);
13211335
}
13221336

1337+
private static ClientStats createClientStats(
1338+
@Nullable Instant requestSendTimeStamp,
1339+
long queuedRequestsCount,
1340+
@Nullable ClientStats.WindowStats windowStats) {
1341+
ClientStats.Builder clientStatsBuilder = ClientStats.newBuilder();
1342+
ClientStats.RequestStats requestStats =
1343+
createRequestStats(requestSendTimeStamp, queuedRequestsCount);
1344+
if (!requestStats.equals(ClientStats.RequestStats.getDefaultInstance())) {
1345+
clientStatsBuilder.setRequestStats(requestStats);
1346+
}
1347+
if (windowStats != null && !windowStats.equals(ClientStats.WindowStats.getDefaultInstance())) {
1348+
clientStatsBuilder.setWindowStats(windowStats);
1349+
}
1350+
return clientStatsBuilder.build();
1351+
}
1352+
1353+
private static ClientStats.RequestStats createRequestStats(
1354+
@Nullable Instant requestSendTimeStamp, long queuedRequestsCount) {
1355+
ClientStats.RequestStats.Builder requestStatsBuilder = ClientStats.RequestStats.newBuilder();
1356+
if (requestSendTimeStamp != null) {
1357+
requestStatsBuilder.setSendTimeMillis(requestSendTimeStamp.toEpochMilli());
1358+
}
1359+
if (queuedRequestsCount > 0) {
1360+
requestStatsBuilder.setQueuedRequestsCount(queuedRequestsCount);
1361+
}
1362+
return requestStatsBuilder.build();
1363+
}
1364+
1365+
@GuardedBy("lock")
1366+
private ClientStats.WindowStats createWindowStats() {
1367+
ClientStats.WindowStats.Builder windowStatsBuilder = ClientStats.WindowStats.newBuilder();
1368+
if (healthCheckMetrics.windowedMilliLatencyMax > 0) {
1369+
windowStatsBuilder.setMaxResponseLatencyMillis(healthCheckMetrics.windowedMilliLatencyMax);
1370+
}
1371+
long avgResponseLatencyMillis =
1372+
healthCheckMetrics.windowedResponsesAcked > 0
1373+
? healthCheckMetrics.windowedMilliLatencySum / healthCheckMetrics.windowedResponsesAcked
1374+
: 0;
1375+
if (avgResponseLatencyMillis > 0) {
1376+
windowStatsBuilder.setAvgResponseLatencyMillis(avgResponseLatencyMillis);
1377+
}
1378+
if (healthCheckMetrics.windowedMilliResponseWaitTimeMax > 0) {
1379+
windowStatsBuilder.setLongestWaitNoResponseMillis(
1380+
healthCheckMetrics.windowedMilliResponseWaitTimeMax);
1381+
}
1382+
if (healthCheckMetrics.windowedRequestsSent > 0) {
1383+
windowStatsBuilder.setRequestsSentCount(healthCheckMetrics.windowedRequestsSent);
1384+
}
1385+
if (healthCheckMetrics.windowedResponsesAcked > 0) {
1386+
windowStatsBuilder.setResponsesReceivedCount(healthCheckMetrics.windowedResponsesAcked);
1387+
}
1388+
if (healthCheckMetrics.windowedRequestsSentBytes > 0) {
1389+
windowStatsBuilder.setBytesSentCount(healthCheckMetrics.windowedRequestsSentBytes);
1390+
}
1391+
long windowStartTimeEpochMillis = healthCheckMetrics.healthCheckTimeStamp.toEpochMilli();
1392+
windowStatsBuilder.setWindowStartTimeEpochMillis(windowStartTimeEpochMillis);
1393+
long windowMillis = Instant.now().toEpochMilli() - windowStartTimeEpochMillis;
1394+
if (windowMillis >= 0) {
1395+
windowStatsBuilder.setWindowMillis(windowMillis);
1396+
}
1397+
return windowStatsBuilder.build();
1398+
}
1399+
13231400
@Nullable
13241401
private AppendRowsSchema getSchema(AppendRowsRequest request) {
13251402
if (request.hasProtoRows()) {

‎java-bigquerystorage/google-cloud-bigquerystorage/src/test/java/com/google/cloud/bigquery/storage/v1/ConnectionWorkerTest.java‎

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1612,6 +1612,17 @@ void testInflightRetryCountHealthMetricExactlyOnce() throws Exception {
16121612
assertTrue(healthCheckFields.responseCodes.containsKey(Status.Code.OK.value()));
16131613
assertEquals(3, healthCheckFields.responseCodes.get(Status.Code.OK.value()));
16141614
assertEquals("projects/p1/datasets/d1/tables/t1/streams/s1", healthCheckFields.streamName);
1615+
1616+
AppendRowsRequest secondRequest = testBigQueryWrite.getAppendRequests().get(1);
1617+
assertThat(secondRequest.getClientStats().getRequestStats().hasQueuedRequestsCount()).isTrue();
1618+
assertThat(secondRequest.getClientStats().getRequestStats().getQueuedRequestsCount())
1619+
.isEqualTo(2);
1620+
// When both requests are retried together in localQueue, only the first request in localQueue
1621+
// should have window_stats populated.
1622+
AppendRowsRequest retryFirstInLocalQueue = testBigQueryWrite.getAppendRequests().get(2);
1623+
AppendRowsRequest retrySecondInLocalQueue = testBigQueryWrite.getAppendRequests().get(3);
1624+
assertThat(retryFirstInLocalQueue.getClientStats().hasWindowStats()).isTrue();
1625+
assertThat(retrySecondInLocalQueue.getClientStats().hasWindowStats()).isFalse();
16151626
}
16161627

16171628
@Test
@@ -1722,6 +1733,114 @@ void testEarliestSendTime_retryScenario() throws Exception {
17221733

17231734
// Once exhausted, state should be null
17241735
assertThat(connectionWorker.getEarliestSendTime()).isNull();
1736+
1737+
AppendRowsRequest initialRequest = testBigQueryWrite.getAppendRequests().get(0);
1738+
AppendRowsRequest retryRequest = testBigQueryWrite.getAppendRequests().get(1);
1739+
assertThat(initialRequest.getClientStats().getRequestStats().hasSendTimeMillis()).isTrue();
1740+
assertThat(retryRequest.getClientStats().getRequestStats().hasSendTimeMillis()).isTrue();
1741+
assertThat(retryRequest.getClientStats().getRequestStats().getSendTimeMillis())
1742+
.isAtLeast(initialRequest.getClientStats().getRequestStats().getSendTimeMillis());
1743+
assertThat(initialRequest.getClientStats().getRequestStats().hasQueuedRequestsCount())
1744+
.isTrue();
1745+
assertThat(initialRequest.getClientStats().getRequestStats().getQueuedRequestsCount())
1746+
.isEqualTo(1);
1747+
assertThat(retryRequest.getClientStats().getRequestStats().hasQueuedRequestsCount()).isTrue();
1748+
assertThat(retryRequest.getClientStats().getRequestStats().getQueuedRequestsCount())
1749+
.isEqualTo(1);
1750+
}
1751+
}
1752+
1753+
@Test
1754+
void testWindowStats() throws Exception {
1755+
long beforeCreationMillis = Instant.now().toEpochMilli();
1756+
ProtoSchema schema1 = createProtoSchema("foo");
1757+
try (StreamWriter sw1 =
1758+
StreamWriter.newBuilder(TEST_STREAM_1, client)
1759+
.setLocation("us")
1760+
.setWriterSchema(schema1)
1761+
.build();
1762+
ConnectionWorker connectionWorker = createConnectionWorker()) {
1763+
long afterCreationMillis = Instant.now().toEpochMilli();
1764+
int msecResponseDelay = 300;
1765+
testBigQueryWrite.setResponseSleep(Duration.ofMillis(msecResponseDelay));
1766+
testBigQueryWrite.addResponse(createAppendResponse(0));
1767+
testBigQueryWrite.addResponse(createAppendResponse(1));
1768+
testBigQueryWrite.addResponse(createAppendResponse(2));
1769+
1770+
// First request: no prior responses or wait times have been recorded yet.
1771+
ApiFuture<AppendRowsResponse> future1 =
1772+
sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"0"}), 0);
1773+
future1.get();
1774+
1775+
AppendRowsRequest firstRequest = testBigQueryWrite.getAppendRequests().get(0);
1776+
long sizePerRequest = firstRequest.getProtoRows().getSerializedSize();
1777+
assertThat(firstRequest.getClientStats().hasWindowStats()).isTrue();
1778+
ClientStats.WindowStats firstWindowStats = firstRequest.getClientStats().getWindowStats();
1779+
assertThat(firstWindowStats.hasMaxResponseLatencyMillis()).isFalse();
1780+
assertThat(firstWindowStats.hasAvgResponseLatencyMillis()).isFalse();
1781+
assertThat(firstWindowStats.hasLongestWaitNoResponseMillis()).isFalse();
1782+
assertThat(firstWindowStats.hasRequestsSentCount()).isTrue();
1783+
assertThat(firstWindowStats.getRequestsSentCount()).isEqualTo(1);
1784+
assertThat(firstWindowStats.hasResponsesReceivedCount()).isFalse();
1785+
assertThat(firstWindowStats.hasBytesSentCount()).isTrue();
1786+
assertThat(firstWindowStats.getBytesSentCount()).isEqualTo(sizePerRequest);
1787+
assertThat(firstWindowStats.hasWindowStartTimeEpochMillis()).isTrue();
1788+
assertThat(firstWindowStats.getWindowStartTimeEpochMillis()).isAtLeast(beforeCreationMillis);
1789+
assertThat(firstWindowStats.getWindowStartTimeEpochMillis()).isAtMost(afterCreationMillis);
1790+
assertThat(firstWindowStats.hasWindowMillis()).isTrue();
1791+
assertThat(firstWindowStats.getWindowMillis()).isAtLeast(0L);
1792+
1793+
// Second request: sent after firstRequest waited in-flight (>100ms) and received a response
1794+
// (>=300ms latency), so all expected WindowStats fields are now non-zero and populated.
1795+
ApiFuture<AppendRowsResponse> future2 =
1796+
sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"1"}), 1);
1797+
future2.get();
1798+
1799+
AppendRowsRequest secondRequest = testBigQueryWrite.getAppendRequests().get(1);
1800+
assertThat(secondRequest.getClientStats().hasWindowStats()).isTrue();
1801+
ClientStats.WindowStats secondWindowStats = secondRequest.getClientStats().getWindowStats();
1802+
assertThat(secondWindowStats.hasMaxResponseLatencyMillis()).isTrue();
1803+
assertThat(secondWindowStats.getMaxResponseLatencyMillis()).isAtLeast(msecResponseDelay);
1804+
assertThat(secondWindowStats.hasAvgResponseLatencyMillis()).isTrue();
1805+
assertThat(secondWindowStats.getAvgResponseLatencyMillis()).isAtLeast(msecResponseDelay);
1806+
assertThat(secondWindowStats.hasLongestWaitNoResponseMillis()).isTrue();
1807+
assertThat(secondWindowStats.getLongestWaitNoResponseMillis()).isGreaterThan(0L);
1808+
assertThat(secondWindowStats.hasRequestsSentCount()).isTrue();
1809+
assertThat(secondWindowStats.getRequestsSentCount()).isEqualTo(2);
1810+
assertThat(secondWindowStats.hasResponsesReceivedCount()).isTrue();
1811+
assertThat(secondWindowStats.getResponsesReceivedCount()).isEqualTo(1);
1812+
assertThat(secondWindowStats.hasBytesSentCount()).isTrue();
1813+
assertThat(secondWindowStats.getBytesSentCount()).isEqualTo(2 * sizePerRequest);
1814+
assertThat(secondWindowStats.hasWindowStartTimeEpochMillis()).isTrue();
1815+
assertThat(secondWindowStats.getWindowStartTimeEpochMillis())
1816+
.isEqualTo(firstWindowStats.getWindowStartTimeEpochMillis());
1817+
assertThat(secondWindowStats.hasWindowMillis()).isTrue();
1818+
assertThat(secondWindowStats.getWindowMillis()).isAtLeast(msecResponseDelay);
1819+
1820+
// Trigger a health check window rollover and verify WindowStats resets for the new window.
1821+
connectionWorker.setTestOnlyHealthCheckInterval(Duration.ofMillis(100));
1822+
Thread.sleep(250);
1823+
1824+
ApiFuture<AppendRowsResponse> future3 =
1825+
sendTestMessage(connectionWorker, sw1, createFooProtoRows(new String[] {"2"}), 2);
1826+
future3.get();
1827+
1828+
AppendRowsRequest thirdRequest = testBigQueryWrite.getAppendRequests().get(2);
1829+
assertThat(thirdRequest.getClientStats().hasWindowStats()).isTrue();
1830+
ClientStats.WindowStats thirdWindowStats = thirdRequest.getClientStats().getWindowStats();
1831+
assertThat(thirdWindowStats.hasMaxResponseLatencyMillis()).isFalse();
1832+
assertThat(thirdWindowStats.hasAvgResponseLatencyMillis()).isFalse();
1833+
assertThat(thirdWindowStats.hasLongestWaitNoResponseMillis()).isFalse();
1834+
assertThat(thirdWindowStats.hasRequestsSentCount()).isTrue();
1835+
assertThat(thirdWindowStats.getRequestsSentCount()).isEqualTo(1);
1836+
assertThat(thirdWindowStats.hasResponsesReceivedCount()).isFalse();
1837+
assertThat(thirdWindowStats.hasBytesSentCount()).isTrue();
1838+
assertThat(thirdWindowStats.getBytesSentCount()).isEqualTo(sizePerRequest);
1839+
assertThat(thirdWindowStats.hasWindowStartTimeEpochMillis()).isTrue();
1840+
assertThat(thirdWindowStats.getWindowStartTimeEpochMillis())
1841+
.isGreaterThan(secondWindowStats.getWindowStartTimeEpochMillis());
1842+
assertThat(thirdWindowStats.hasWindowMillis()).isTrue();
1843+
assertThat(thirdWindowStats.getWindowMillis()).isAtLeast(0L);
17251844
}
17261845
}
17271846
}

0 commit comments

Comments
 (0)