Browse Source
Adds page-level byte-size tracking in TbResultSet.allRows() using ExecutionInfo.getResponseSizeInBytes() to fail early when accumulated result set size exceeds the configurable limit (default 50MB). Also fixes pre-existing bugs where onFailure callbacks in CassandraBaseTimeseriesDao only logged errors but never completed the future, causing callers to hang indefinitely on failures. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>pull/15058/head
Viacheslav Klimov
3 months ago
9 changed files with 295 additions and 42 deletions
@ -0,0 +1,32 @@ |
|||
/** |
|||
* Copyright © 2016-2026 The Thingsboard Authors |
|||
* |
|||
* Licensed under the Apache License, Version 2.0 (the "License"); |
|||
* you may not use this file except in compliance with the License. |
|||
* You may obtain a copy of the License at |
|||
* |
|||
* http://www.apache.org/licenses/LICENSE-2.0
|
|||
* |
|||
* Unless required by applicable law or agreed to in writing, software |
|||
* distributed under the License is distributed on an "AS IS" BASIS, |
|||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|||
* See the License for the specific language governing permissions and |
|||
* limitations under the License. |
|||
*/ |
|||
package org.thingsboard.server.dao.nosql; |
|||
|
|||
import lombok.Getter; |
|||
|
|||
@Getter |
|||
public class ResultSetSizeLimitExceededException extends IllegalArgumentException { |
|||
|
|||
private final long limitBytes; |
|||
private final long actualBytes; |
|||
|
|||
public ResultSetSizeLimitExceededException(long limitBytes, long actualBytes) { |
|||
super("Result set size exceeds the maximum allowed limit. Please narrow your query"); |
|||
this.limitBytes = limitBytes; |
|||
this.actualBytes = actualBytes; |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,140 @@ |
|||
/** |
|||
* Copyright © 2016-2026 The Thingsboard Authors |
|||
* |
|||
* Licensed under the Apache License, Version 2.0 (the "License"); |
|||
* you may not use this file except in compliance with the License. |
|||
* You may obtain a copy of the License at |
|||
* |
|||
* http://www.apache.org/licenses/LICENSE-2.0
|
|||
* |
|||
* Unless required by applicable law or agreed to in writing, software |
|||
* distributed under the License is distributed on an "AS IS" BASIS, |
|||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
|||
* See the License for the specific language governing permissions and |
|||
* limitations under the License. |
|||
*/ |
|||
package org.thingsboard.server.dao.nosql; |
|||
|
|||
import com.datastax.oss.driver.api.core.cql.AsyncResultSet; |
|||
import com.datastax.oss.driver.api.core.cql.ColumnDefinitions; |
|||
import com.datastax.oss.driver.api.core.cql.ExecutionInfo; |
|||
import com.datastax.oss.driver.api.core.cql.Row; |
|||
import com.datastax.oss.driver.api.core.cql.Statement; |
|||
import com.google.common.util.concurrent.ListenableFuture; |
|||
import com.google.common.util.concurrent.MoreExecutors; |
|||
import com.google.common.util.concurrent.SettableFuture; |
|||
import org.junit.jupiter.api.Test; |
|||
|
|||
import java.nio.ByteBuffer; |
|||
import java.util.List; |
|||
import java.util.concurrent.ExecutionException; |
|||
import java.util.function.Function; |
|||
|
|||
import static org.assertj.core.api.Assertions.assertThat; |
|||
import static org.assertj.core.api.Assertions.assertThatThrownBy; |
|||
import static org.mockito.Mockito.doReturn; |
|||
import static org.mockito.Mockito.mock; |
|||
import static org.mockito.Mockito.when; |
|||
|
|||
class TbResultSetTest { |
|||
|
|||
@Test |
|||
void allRows_withinLimit_returnsAllRows() throws Exception { |
|||
Row row = mock(Row.class); |
|||
AsyncResultSet asyncResultSet = createMockResultSet(List.of(row), false, 1000); |
|||
Statement<?> statement = mock(Statement.class); |
|||
|
|||
TbResultSet tbResultSet = new TbResultSet(statement, asyncResultSet, s -> null); |
|||
ListenableFuture<List<Row>> future = tbResultSet.allRows(MoreExecutors.directExecutor(), 5000); |
|||
|
|||
List<Row> result = future.get(); |
|||
assertThat(result).hasSize(1); |
|||
assertThat(result.get(0)).isSameAs(row); |
|||
} |
|||
|
|||
@Test |
|||
void allRows_exceedsLimitOnFirstPage_failsWithException() { |
|||
Row row = mock(Row.class); |
|||
AsyncResultSet asyncResultSet = createMockResultSet(List.of(row), false, 6000); |
|||
Statement<?> statement = mock(Statement.class); |
|||
|
|||
TbResultSet tbResultSet = new TbResultSet(statement, asyncResultSet, s -> null); |
|||
ListenableFuture<List<Row>> future = tbResultSet.allRows(MoreExecutors.directExecutor(), 5000); |
|||
|
|||
assertThatThrownBy(future::get) |
|||
.isInstanceOf(ExecutionException.class) |
|||
.hasCauseInstanceOf(ResultSetSizeLimitExceededException.class); |
|||
} |
|||
|
|||
@Test |
|||
void allRows_exceedsLimitOnSecondPage_failsAfterSecondPage() { |
|||
Row row1 = mock(Row.class); |
|||
Row row2 = mock(Row.class); |
|||
Statement<?> statement = mock(Statement.class); |
|||
doReturn(statement).when(statement).setPagingState((ByteBuffer) null); |
|||
|
|||
AsyncResultSet page2 = createMockResultSet(List.of(row2), false, 3000); |
|||
TbResultSet tbResultSetPage2 = new TbResultSet(statement, page2, s -> null); |
|||
SettableFuture<TbResultSet> page2Future = SettableFuture.create(); |
|||
page2Future.set(tbResultSetPage2); |
|||
TbResultSetFuture tbPage2Future = new TbResultSetFuture(page2Future); |
|||
|
|||
ExecutionInfo page1ExecInfo = mock(ExecutionInfo.class); |
|||
when(page1ExecInfo.getResponseSizeInBytes()).thenReturn(3000); |
|||
when(page1ExecInfo.getPagingState()).thenReturn(null); |
|||
|
|||
AsyncResultSet page1 = createMockResultSet(List.of(row1), true, 3000); |
|||
when(page1.getExecutionInfo()).thenReturn(page1ExecInfo); |
|||
|
|||
Function<Statement, TbResultSetFuture> executeAsync = s -> tbPage2Future; |
|||
TbResultSet tbResultSet = new TbResultSet(statement, page1, executeAsync); |
|||
ListenableFuture<List<Row>> future = tbResultSet.allRows(MoreExecutors.directExecutor(), 5000); |
|||
|
|||
assertThatThrownBy(future::get) |
|||
.isInstanceOf(ExecutionException.class) |
|||
.hasCauseInstanceOf(ResultSetSizeLimitExceededException.class); |
|||
} |
|||
|
|||
@Test |
|||
void allRows_unlimitedWithZero_returnsAllRowsRegardlessOfSize() throws Exception { |
|||
Row row = mock(Row.class); |
|||
AsyncResultSet asyncResultSet = createMockResultSet(List.of(row), false, 999999); |
|||
Statement<?> statement = mock(Statement.class); |
|||
|
|||
TbResultSet tbResultSet = new TbResultSet(statement, asyncResultSet, s -> null); |
|||
ListenableFuture<List<Row>> future = tbResultSet.allRows(MoreExecutors.directExecutor(), 0); |
|||
|
|||
List<Row> result = future.get(); |
|||
assertThat(result).hasSize(1); |
|||
} |
|||
|
|||
@Test |
|||
void allRows_noLimitOverload_returnsAllRows() throws Exception { |
|||
Row row = mock(Row.class); |
|||
AsyncResultSet asyncResultSet = createMockResultSet(List.of(row), false, 999999); |
|||
Statement<?> statement = mock(Statement.class); |
|||
|
|||
TbResultSet tbResultSet = new TbResultSet(statement, asyncResultSet, s -> null); |
|||
ListenableFuture<List<Row>> future = tbResultSet.allRows(MoreExecutors.directExecutor()); |
|||
|
|||
List<Row> result = future.get(); |
|||
assertThat(result).hasSize(1); |
|||
} |
|||
|
|||
private AsyncResultSet createMockResultSet(List<Row> rows, boolean hasMorePages, int responseSizeInBytes) { |
|||
AsyncResultSet resultSet = mock(AsyncResultSet.class); |
|||
ExecutionInfo executionInfo = mock(ExecutionInfo.class); |
|||
ColumnDefinitions columnDefs = mock(ColumnDefinitions.class); |
|||
|
|||
when(executionInfo.getResponseSizeInBytes()).thenReturn(responseSizeInBytes); |
|||
when(executionInfo.getPagingState()).thenReturn(null); |
|||
when(resultSet.getExecutionInfo()).thenReturn(executionInfo); |
|||
when(resultSet.getColumnDefinitions()).thenReturn(columnDefs); |
|||
when(resultSet.currentPage()).thenReturn(rows); |
|||
when(resultSet.hasMorePages()).thenReturn(hasMorePages); |
|||
when(resultSet.remaining()).thenReturn(rows.size()); |
|||
|
|||
return resultSet; |
|||
} |
|||
|
|||
} |
|||
Loading…
Reference in new issue