|
21 | 21 | import com.google.api.core.InternalApi; |
22 | 22 | import com.google.api.gax.grpc.GrpcCallContext; |
23 | 23 | import com.google.api.gax.rpc.ApiCallContext; |
| 24 | +import com.google.api.gax.rpc.ResponseObserver; |
24 | 25 | import com.google.api.gax.rpc.ServerStream; |
25 | 26 | import com.google.api.gax.rpc.StateCheckingResponseObserver; |
26 | 27 | import com.google.api.gax.rpc.StreamController; |
|
43 | 44 | import com.google.cloud.bigtable.metrics.Timer; |
44 | 45 | import com.google.cloud.bigtable.metrics.Timer.Context; |
45 | 46 | import com.google.common.util.concurrent.MoreExecutors; |
| 47 | +import com.google.common.util.concurrent.SettableFuture; |
46 | 48 | import com.google.protobuf.ByteString; |
47 | 49 | import io.grpc.CallOptions; |
48 | 50 | import io.grpc.Deadline; |
49 | 51 | import io.grpc.stub.StreamObserver; |
| 52 | +import java.util.ArrayDeque; |
| 53 | +import java.util.ArrayList; |
50 | 54 | import java.util.Iterator; |
51 | 55 | import java.util.List; |
| 56 | +import java.util.Queue; |
| 57 | +import java.util.concurrent.ExecutionException; |
| 58 | +import java.util.concurrent.Future; |
52 | 59 | import java.util.concurrent.TimeUnit; |
53 | 60 | import javax.annotation.Nullable; |
54 | 61 | import org.apache.hadoop.hbase.client.AbstractClientScanner; |
@@ -134,6 +141,12 @@ public Result apply(Row row) { |
134 | 141 | MoreExecutors.directExecutor()); |
135 | 142 | } |
136 | 143 |
|
| 144 | + @Override |
| 145 | + public ResultScanner readRows(Query.QueryPaginator paginator, long maxSegmentByteSize) { |
| 146 | + return new PaginatedRowResultScanner( |
| 147 | + paginator, delegate, maxSegmentByteSize, this.createScanCallContext()); |
| 148 | + } |
| 149 | + |
137 | 150 | @Override |
138 | 151 | public ResultScanner readRows(Query request) { |
139 | 152 | return new RowResultScanner( |
@@ -228,6 +241,151 @@ protected void onCompleteImpl() { |
228 | 241 | } |
229 | 242 | } |
230 | 243 |
|
| 244 | + /** |
| 245 | + * wraps {@link ServerStream} onto HBase {@link ResultScanner}. {@link PaginatedRowResultScanner} |
| 246 | + * gets a paginator and a {@link Query.QueryPaginator} used to get a {@link ServerStream}<{@link |
| 247 | + * Result}> using said paginator to iterate over pages of rows. The {@link Query.QueryPaginator} |
| 248 | + * pageSize property indicates the size of each page in every API call. A cache of a maximum size |
| 249 | + * of 1.1*pageSize and a minimum of 0.1*pageSize is held at all times. In order to avoid OOM |
| 250 | + * exceptions, there is a limit for the total byte size held in cache. |
| 251 | + */ |
| 252 | + static class PaginatedRowResultScanner extends AbstractClientScanner { |
| 253 | + // Percentage of max number of rows allowed in the buffer |
| 254 | + private static final double WATERMARK_PERCENTAGE = .1; |
| 255 | + private static final RowResultAdapter RESULT_ADAPTER = new RowResultAdapter(); |
| 256 | + |
| 257 | + private final Meter scannerResultMeter = |
| 258 | + BigtableClientMetrics.meter(BigtableClientMetrics.MetricLevel.Info, "scanner.results"); |
| 259 | + private final Timer scannerResultTimer = |
| 260 | + BigtableClientMetrics.timer( |
| 261 | + BigtableClientMetrics.MetricLevel.Debug, "scanner.results.latency"); |
| 262 | + |
| 263 | + private ByteString lastSeenRowKey = ByteString.EMPTY; |
| 264 | + private Boolean hasMore = true; |
| 265 | + private final Queue<Result> buffer; |
| 266 | + private final Query.QueryPaginator paginator; |
| 267 | + private final int refillSegmentWaterMark; |
| 268 | + |
| 269 | + private final BigtableDataClient dataClient; |
| 270 | + |
| 271 | + private final long maxSegmentByteSize; |
| 272 | + |
| 273 | + private long currentByteSize = 0; |
| 274 | + |
| 275 | + private @Nullable Future<List<Result>> future; |
| 276 | + private GrpcCallContext scanCallContext; |
| 277 | + |
| 278 | + PaginatedRowResultScanner( |
| 279 | + Query.QueryPaginator paginator, |
| 280 | + BigtableDataClient dataClient, |
| 281 | + long maxSegmentByteSize, |
| 282 | + GrpcCallContext scanCallContext) { |
| 283 | + this.maxSegmentByteSize = maxSegmentByteSize; |
| 284 | + |
| 285 | + this.paginator = paginator; |
| 286 | + this.dataClient = dataClient; |
| 287 | + this.buffer = new ArrayDeque<>(); |
| 288 | + this.refillSegmentWaterMark = |
| 289 | + (int) Math.max(1, paginator.getPageSize() * WATERMARK_PERCENTAGE); |
| 290 | + this.scanCallContext = scanCallContext; |
| 291 | + this.future = fetchNextSegment(); |
| 292 | + } |
| 293 | + |
| 294 | + @Override |
| 295 | + public Result next() { |
| 296 | + try (Context ignored = scannerResultTimer.time()) { |
| 297 | + if (this.future != null && this.future.isDone()) { |
| 298 | + this.consumeReadRowsFuture(); |
| 299 | + } |
| 300 | + if (this.buffer.size() < this.refillSegmentWaterMark && this.future == null && hasMore) { |
| 301 | + future = fetchNextSegment(); |
| 302 | + } |
| 303 | + if (this.buffer.isEmpty() && this.future != null) { |
| 304 | + this.consumeReadRowsFuture(); |
| 305 | + } |
| 306 | + Result result = this.buffer.poll(); |
| 307 | + if (result != null) { |
| 308 | + scannerResultMeter.mark(); |
| 309 | + currentByteSize -= Result.getTotalSizeOfCells(result); |
| 310 | + } |
| 311 | + return result; |
| 312 | + } |
| 313 | + } |
| 314 | + |
| 315 | + @Override |
| 316 | + public void close() { |
| 317 | + if (this.future != null) { |
| 318 | + this.future.cancel(true); |
| 319 | + } |
| 320 | + } |
| 321 | + |
| 322 | + public boolean renewLease() { |
| 323 | + return true; |
| 324 | + } |
| 325 | + |
| 326 | + private Future<List<Result>> fetchNextSegment() { |
| 327 | + SettableFuture<List<Result>> resultsFuture = SettableFuture.create(); |
| 328 | + |
| 329 | + dataClient |
| 330 | + .readRowsCallable(RESULT_ADAPTER) |
| 331 | + .call( |
| 332 | + paginator.getNextQuery(), |
| 333 | + new ResponseObserver<Result>() { |
| 334 | + private StreamController controller; |
| 335 | + List<Result> results = new ArrayList(); |
| 336 | + |
| 337 | + @Override |
| 338 | + public void onStart(StreamController controller) { |
| 339 | + this.controller = controller; |
| 340 | + } |
| 341 | + |
| 342 | + @Override |
| 343 | + public void onResponse(Result result) { |
| 344 | + // calculate size of the response |
| 345 | + currentByteSize += Result.getTotalSizeOfCells(result); |
| 346 | + results.add(result); |
| 347 | + if (result != null && result.rawCells() != null) { |
| 348 | + lastSeenRowKey = RESULT_ADAPTER.getKey(result); |
| 349 | + } |
| 350 | + |
| 351 | + if (currentByteSize > maxSegmentByteSize) { |
| 352 | + controller.cancel(); |
| 353 | + return; |
| 354 | + } |
| 355 | + } |
| 356 | + |
| 357 | + @Override |
| 358 | + public void onError(Throwable t) { |
| 359 | + if (currentByteSize > maxSegmentByteSize) { |
| 360 | + onComplete(); |
| 361 | + } else { |
| 362 | + resultsFuture.setException(t); |
| 363 | + } |
| 364 | + } |
| 365 | + |
| 366 | + @Override |
| 367 | + public void onComplete() { |
| 368 | + resultsFuture.set(results); |
| 369 | + } |
| 370 | + }, |
| 371 | + this.scanCallContext); |
| 372 | + return resultsFuture; |
| 373 | + } |
| 374 | + |
| 375 | + private void consumeReadRowsFuture() { |
| 376 | + try { |
| 377 | + List<Result> results = this.future.get(); |
| 378 | + this.buffer.addAll(results); |
| 379 | + this.hasMore = this.paginator.advance(this.lastSeenRowKey); |
| 380 | + this.future = null; |
| 381 | + } catch (InterruptedException e) { |
| 382 | + Thread.currentThread().interrupt(); |
| 383 | + } catch (ExecutionException e) { |
| 384 | + // Do nothing. |
| 385 | + } |
| 386 | + } |
| 387 | + } |
| 388 | + |
231 | 389 | /** wraps {@link ServerStream} onto HBase {@link ResultScanner}. */ |
232 | 390 | private static class RowResultScanner extends AbstractClientScanner { |
233 | 391 |
|
@@ -264,7 +422,7 @@ public void close() { |
264 | 422 | } |
265 | 423 |
|
266 | 424 | public boolean renewLease() { |
267 | | - throw new UnsupportedOperationException("renewLease"); |
| 425 | + return true; |
268 | 426 | } |
269 | 427 | } |
270 | 428 | } |
0 commit comments