Source:
<code>public boolean load(List<UpdateRequest> requests) { UnloadedIndexKeyRecordSet recordSetForLookupSource = new UnloadedIndexKeyRecordSet(pipelineContext.getSession(), indexSnapshotReference.get(), lookupSourceInputChannels, indexTypes, requests, joinCompiler); try (Driver driver = driverFactory.createDriver(pipelineContext.addDriverContext())) { PlanNodeId sourcePlanNodeId = driverFactory.getSourceId().get(); ScheduledSplit split = new ScheduledSplit(0, sourcePlanNodeId, new Split(INDEX_CONNECTOR_ID, new ConnectorTransactionHandle() {}, new IndexSplit(recordSetForLookupSource))); driver.updateSource(new TaskSource(sourcePlanNodeId, ImmutableSet.of(split), true)); while (!driver.isFinished()) { ListenableFuture<?> process = driver.process(); long start = System.currentTimeMillis(); <START> while (!process.isDone() && System.currentTimeMillis() - start > timeout.toMillis()) { <END> try { sleep(10); } catch (InterruptedException e) { break; } } if (!process.isDone()) { throw new PrestoException(INDEX_LOADER_TIMEOUT, format("Exceeded the time limit of %s loading indexes for index join", timeout)); } } } if (indexSnapshotBuilder.isMemoryExceeded()) { clearCachedData(); return false; } UnloadedIndexKeyRecordSet indexKeysRecordSet = (lookupSourceInputChannels.equals(allInputChannels)) ? recordSetForLookupSource : new UnloadedIndexKeyRecordSet(pipelineContext.getSession(), indexSnapshotReference.get(), allInputChannels, indexTypes, requests, joinCompiler); IndexSnapshot newValue = indexSnapshotBuilder.createIndexSnapshot(indexKeysRecordSet); if (newValue == null) { clearCachedData(); return false; } indexSnapshotReference.set(newValue); for (UpdateRequest request : requests) { request.finished(newValue); } return true; }</code><technical_language>process.get(timeout.toMillis(), TimeUnit.MILLISECONDS); + some try/catching handling maybe</technical_language>

Target:
public boolean load(List<UpdateRequest> requests) { UnloadedIndexKeyRecordSet recordSetForLookupSource = new UnloadedIndexKeyRecordSet(pipelineContext.getSession(), indexSnapshotReference.get(), lookupSourceInputChannels, indexTypes, requests, joinCompiler); try (Driver driver = driverFactory.createDriver(pipelineContext.addDriverContext())) { PlanNodeId sourcePlanNodeId = driverFactory.getSourceId().get(); ScheduledSplit split = new ScheduledSplit(0, sourcePlanNodeId, new Split(INDEX_CONNECTOR_ID, new ConnectorTransactionHandle() {}, new IndexSplit(recordSetForLookupSource))); driver.updateSource(new TaskSource(sourcePlanNodeId, ImmutableSet.of(split), true)); while (!driver.isFinished()) { ListenableFuture<?> process = driver.process(); try { process.get(timeout.toMillis(), MILLISECONDS); } catch (TimeoutException e) { throw new PrestoException(INDEX_LOADER_TIMEOUT, format("Exceeded the time limit of %s loading indexes for index join", timeout)); } catch (ExecutionException e) { throw new PrestoException(GENERIC_INTERNAL_ERROR, "Error loading index for join", e); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new PrestoException(GENERIC_INTERNAL_ERROR, "Error loading index for join", e); } } } if (indexSnapshotBuilder.isMemoryExceeded()) { clearCachedData(); return false; } UnloadedIndexKeyRecordSet indexKeysRecordSet = (lookupSourceInputChannels.equals(allInputChannels)) ? recordSetForLookupSource : new UnloadedIndexKeyRecordSet(pipelineContext.getSession(), indexSnapshotReference.get(), allInputChannels, indexTypes, requests, joinCompiler); IndexSnapshot newValue = indexSnapshotBuilder.createIndexSnapshot(indexKeysRecordSet); if (newValue == null) { clearCachedData(); return false; } indexSnapshotReference.set(newValue); for (UpdateRequest request : requests) { request.finished(newValue); } return true; }