ExecutionGraphCache implements Closeable { public CompletableFuture<AccessExecutionGraph> getExecutionGraph(JobID jobId, RestfulGateway restfulGateway) { Preconditions.checkState(running, "ExecutionGraphCache is no longer running"); while (true) { final ExecutionGraphEntry oldEntry = cachedExecutionGraphs.get(jobId); final long currentTime = System.currentTimeMillis(); if (oldEntry != null) { if (currentTime < oldEntry.getTTL()) { if (oldEntry.getExecutionGraphFuture().isDone() && !oldEntry.getExecutionGraphFuture().isCompletedExceptionally()) { try { if (oldEntry.getExecutionGraphFuture().get().getState() != JobStatus.SUSPENDED) { return oldEntry.getExecutionGraphFuture(); } } catch (InterruptedException | ExecutionException e) { throw new RuntimeException("Could not retrieve ExecutionGraph from the orderly completed future. This should never happen.", e); } } else if (!oldEntry.getExecutionGraphFuture().isDone()) { return oldEntry.getExecutionGraphFuture(); } } } final ExecutionGraphEntry newEntry = new ExecutionGraphEntry(currentTime + timeToLive.toMilliseconds()); final boolean successfulUpdate; if (oldEntry == null) { successfulUpdate = cachedExecutionGraphs.putIfAbsent(jobId, newEntry) == null; } else { successfulUpdate = cachedExecutionGraphs.replace(jobId, oldEntry, newEntry); oldEntry.getExecutionGraphFuture().cancel(false); } if (successfulUpdate) { final CompletableFuture<AccessExecutionGraph> executionGraphFuture = restfulGateway.requestJob(jobId, timeout); executionGraphFuture.whenComplete( (AccessExecutionGraph executionGraph, Throwable throwable) -> { if (throwable != null) { newEntry.getExecutionGraphFuture().completeExceptionally(throwable); cachedExecutionGraphs.remove(jobId, newEntry); } else { newEntry.getExecutionGraphFuture().complete(executionGraph); if (executionGraph.getState() == JobStatus.SUSPENDED) { cachedExecutionGraphs.remove(jobId, newEntry); } } }); if (!running) { cachedExecutionGraphs.remove(jobId, newEntry); } return newEntry.getExecutionGraphFuture(); } } } ExecutionGraphCache(
			Time timeout,
			Time timeToLive); @Override void close(); int size(); CompletableFuture<AccessExecutionGraph> getExecutionGraph(JobID jobId, RestfulGateway restfulGateway); void cleanup();  }
@Test public void testExecutionGraphCaching() throws Exception { final Time timeout = Time.milliseconds(100L); final Time timeToLive = Time.hours(1L); final JobID jobId = new JobID(); final AccessExecutionGraph accessExecutionGraph = mock(AccessExecutionGraph.class); final JobManagerGateway jobManagerGateway = mock(JobManagerGateway.class); when(jobManagerGateway.requestJob(eq(jobId), any(Time.class))).thenReturn(CompletableFuture.completedFuture(accessExecutionGraph)); try (ExecutionGraphCache executionGraphCache = new ExecutionGraphCache(timeout, timeToLive)) { CompletableFuture<AccessExecutionGraph> accessExecutionGraphFuture = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, accessExecutionGraphFuture.get()); CompletableFuture<AccessExecutionGraph> accessExecutionGraphFuture2 = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, accessExecutionGraphFuture2.get()); verify(jobManagerGateway, times(1)).requestJob(eq(jobId), any(Time.class)); } }
@Test public void testExecutionGraphEntryInvalidation() throws Exception { final Time timeout = Time.milliseconds(100L); final Time timeToLive = Time.milliseconds(1L); final JobID jobId = new JobID(); final AccessExecutionGraph accessExecutionGraph = mock(AccessExecutionGraph.class); final JobManagerGateway jobManagerGateway = mock(JobManagerGateway.class); when(jobManagerGateway.requestJob(eq(jobId), any(Time.class))).thenReturn(CompletableFuture.completedFuture(accessExecutionGraph)); try (ExecutionGraphCache executionGraphCache = new ExecutionGraphCache(timeout, timeToLive)) { CompletableFuture<AccessExecutionGraph> executionGraphFuture = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, executionGraphFuture.get()); Thread.sleep(timeToLive.toMilliseconds()); CompletableFuture<AccessExecutionGraph> executionGraphFuture2 = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, executionGraphFuture2.get()); verify(jobManagerGateway, times(2)).requestJob(eq(jobId), any(Time.class)); } }
@Test public void testImmediateCacheInvalidationAfterFailure() throws Exception { final Time timeout = Time.milliseconds(100L); final Time timeToLive = Time.hours(1L); final JobID jobId = new JobID(); final AccessExecutionGraph accessExecutionGraph = mock(AccessExecutionGraph.class); final JobManagerGateway jobManagerGateway = mock(JobManagerGateway.class); when(jobManagerGateway.requestJob(eq(jobId), any(Time.class))).thenReturn( FutureUtils.completedExceptionally(new FlinkJobNotFoundException(jobId)), CompletableFuture.completedFuture(accessExecutionGraph)); try (ExecutionGraphCache executionGraphCache = new ExecutionGraphCache(timeout, timeToLive)) { CompletableFuture<AccessExecutionGraph> executionGraphFuture = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); try { executionGraphFuture.get(); fail("The execution graph future should have been completed exceptionally."); } catch (ExecutionException ee) { assertTrue(ee.getCause() instanceof FlinkException); } CompletableFuture<AccessExecutionGraph> executionGraphFuture2 = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, executionGraphFuture2.get()); } }
@Test public void testConcurrentAccess() throws Exception { final Time timeout = Time.milliseconds(100L); final Time timeToLive = Time.hours(1L); final JobID jobId = new JobID(); final AccessExecutionGraph accessExecutionGraph = mock(AccessExecutionGraph.class); final JobManagerGateway jobManagerGateway = mock(JobManagerGateway.class); when(jobManagerGateway.requestJob(eq(jobId), any(Time.class))).thenReturn(CompletableFuture.completedFuture(accessExecutionGraph)); final int numConcurrentAccesses = 10; final ArrayList<CompletableFuture<AccessExecutionGraph>> executionGraphFutures = new ArrayList<>(numConcurrentAccesses); final ExecutorService executor = java.util.concurrent.Executors.newFixedThreadPool(numConcurrentAccesses); try (ExecutionGraphCache executionGraphCache = new ExecutionGraphCache(timeout, timeToLive)) { for (int i = 0; i < numConcurrentAccesses; i++) { CompletableFuture<AccessExecutionGraph> executionGraphFuture = CompletableFuture .supplyAsync( () -> executionGraphCache.getExecutionGraph(jobId, jobManagerGateway), executor) .thenCompose(Function.identity()); executionGraphFutures.add(executionGraphFuture); } final CompletableFuture<Collection<AccessExecutionGraph>> allExecutionGraphFutures = FutureUtils.combineAll(executionGraphFutures); Collection<AccessExecutionGraph> allExecutionGraphs = allExecutionGraphFutures.get(); for (AccessExecutionGraph executionGraph : allExecutionGraphs) { assertEquals(accessExecutionGraph, executionGraph); } verify(jobManagerGateway, times(1)).requestJob(eq(jobId), any(Time.class)); } finally { ExecutorUtils.gracefulShutdown(5000L, TimeUnit.MILLISECONDS, executor); } }
@Test public void testCacheInvalidationIfSuspended() throws Exception { final Time timeout = Time.milliseconds(100L); final Time timeToLive = Time.hours(1L); final JobID jobId = new JobID(); final AccessExecutionGraph accessExecutionGraph = mock(AccessExecutionGraph.class); final AccessExecutionGraph suspendedExecutionGraph = mock(AccessExecutionGraph.class); when(suspendedExecutionGraph.getState()).thenReturn(JobStatus.SUSPENDED); final JobManagerGateway jobManagerGateway = mock(JobManagerGateway.class); when(jobManagerGateway.requestJob(eq(jobId), any(Time.class))).thenReturn( CompletableFuture.completedFuture(suspendedExecutionGraph), CompletableFuture.completedFuture(accessExecutionGraph)); try (ExecutionGraphCache executionGraphCache = new ExecutionGraphCache(timeout, timeToLive)) { CompletableFuture<AccessExecutionGraph> executionGraphFuture = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(suspendedExecutionGraph, executionGraphFuture.get()); CompletableFuture<AccessExecutionGraph> executionGraphFuture2 = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, executionGraphFuture2.get()); } }
@Test public void testCacheInvalidationIfSwitchToSuspended() throws Exception { final Time timeout = Time.milliseconds(100L); final Time timeToLive = Time.hours(1L); final JobID jobId = new JobID(); final AccessExecutionGraph accessExecutionGraph = mock(AccessExecutionGraph.class); final SuspendableAccessExecutionGraph toBeSuspendedExecutionGraph = new SuspendableAccessExecutionGraph(jobId); final JobManagerGateway jobManagerGateway = mock(JobManagerGateway.class); when(jobManagerGateway.requestJob(eq(jobId), any(Time.class))).thenReturn( CompletableFuture.completedFuture(toBeSuspendedExecutionGraph), CompletableFuture.completedFuture(accessExecutionGraph)); try (ExecutionGraphCache executionGraphCache = new ExecutionGraphCache(timeout, timeToLive)) { CompletableFuture<AccessExecutionGraph> executionGraphFuture = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(toBeSuspendedExecutionGraph, executionGraphFuture.get()); toBeSuspendedExecutionGraph.setJobStatus(JobStatus.SUSPENDED); CompletableFuture<AccessExecutionGraph> executionGraphFuture2 = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, executionGraphFuture2.get()); CompletableFuture<AccessExecutionGraph> executionGraphFuture3 = executionGraphCache.getExecutionGraph(jobId, jobManagerGateway); assertEquals(accessExecutionGraph, executionGraphFuture3.get()); verify(jobManagerGateway, times(2)).requestJob(eq(jobId), any(Time.class)); } }