StateFetcher { public @Nullable <T, SideWindowT extends BoundedWindow> Optional<T> fetchSideInput( final PCollectionView<T> view, final SideWindowT sideWindow, final String stateFamily, SideInputState state, final Supplier<Closeable> scopedReadStateSupplier) { final SideInputId id = new SideInputId(view.getTagInternal(), sideWindow); Callable<SideInputCacheEntry> fetchCallable = () -> { @SuppressWarnings("unchecked") WindowingStrategy<?, SideWindowT> sideWindowStrategy = (WindowingStrategy<?, SideWindowT>) view.getWindowingStrategyInternal(); Coder<SideWindowT> windowCoder = sideWindowStrategy.getWindowFn().windowCoder(); ByteString.Output windowStream = ByteString.newOutput(); windowCoder.encode(sideWindow, windowStream, Coder.Context.OUTER); @SuppressWarnings("unchecked") Windmill.GlobalDataRequest request = Windmill.GlobalDataRequest.newBuilder() .setDataId( Windmill.GlobalDataId.newBuilder() .setTag(view.getTagInternal().getId()) .setVersion(windowStream.toByteString()) .build()) .setStateFamily(stateFamily) .setExistenceWatermarkDeadline( WindmillTimeUtils.harnessToWindmillTimestamp( sideWindowStrategy .getTrigger() .getWatermarkThatGuaranteesFiring(sideWindow))) .build(); Windmill.GlobalData data; try (Closeable scope = scopedReadStateSupplier.get()) { data = server.getSideInputData(request); } bytesRead += data.getSerializedSize(); checkState( SUPPORTED_MATERIALIZATIONS.contains(view.getViewFn().getMaterialization().getUrn()), "Only materializations of type %s supported, received %s", SUPPORTED_MATERIALIZATIONS, view.getViewFn().getMaterialization().getUrn()); Iterable<?> rawData; if (data.getIsReady()) { if (data.getData().size() > 0) { rawData = IterableCoder.of(view.getCoderInternal()) .decode(data.getData().newInput(), Coder.Context.OUTER); } else { rawData = Collections.emptyList(); } switch (view.getViewFn().getMaterialization().getUrn()) { case Materializations.ITERABLE_MATERIALIZATION_URN: { ViewFn<IterableView, T> viewFn = (ViewFn<IterableView, T>) view.getViewFn(); return SideInputCacheEntry.ready( viewFn.apply(() -> rawData), data.getData().size()); } case Materializations.MULTIMAP_MATERIALIZATION_URN: { ViewFn<MultimapView, T> viewFn = (ViewFn<MultimapView, T>) view.getViewFn(); Coder<?> keyCoder = ((KvCoder<?, ?>) view.getCoderInternal()).getKeyCoder(); return SideInputCacheEntry.ready( viewFn.apply( InMemoryMultimapSideInputView.fromIterable(keyCoder, (Iterable) rawData)), data.getData().size()); } default: throw new IllegalStateException( String.format( "Unknown side input materialization format requested '%s'", view.getViewFn().getMaterialization().getUrn())); } } else { return SideInputCacheEntry.notReady(); } }; try { if (state == SideInputState.KNOWN_READY) { SideInputCacheEntry entry = sideInputCache.getIfPresent(id); if (entry == null) { return sideInputCache.get(id, fetchCallable).getValue(); } else if (!entry.isReady()) { synchronized (entry) { SideInputCacheEntry newEntry = sideInputCache.getIfPresent(id); if (newEntry != null && !newEntry.isReady()) { sideInputCache.invalidate(id); } } return sideInputCache.get(id, fetchCallable).getValue(); } else { return entry.getValue(); } } else { return sideInputCache.get(id, fetchCallable).getValue(); } } catch (Exception e) { LOG.error("Fetch failed: ", e); throw new RuntimeException("Exception while fetching side input: ", e); } } StateFetcher(MetricTrackingWindmillServerStub server);  StateFetcher(
      MetricTrackingWindmillServerStub server,
      Cache<SideInputId, SideInputCacheEntry> sideInputCache); StateFetcher byteTrackingView(); long getBytesRead(); @Nullable Optional<T> fetchSideInput(
      final PCollectionView<T> view,
      final SideWindowT sideWindow,
      final String stateFamily,
      SideInputState state,
      final Supplier<Closeable> scopedReadStateSupplier);  }
@Test public void testFetchGlobalDataBasic() throws Exception { StateFetcher fetcher = new StateFetcher(server); ByteString.Output stream = ByteString.newOutput(); ListCoder.of(StringUtf8Coder.of()).encode(Arrays.asList("data"), stream, Coder.Context.OUTER); ByteString encodedIterable = stream.toByteString(); PCollectionView<String> view = TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton()); String tag = view.getTagInternal().getId(); when(server.getSideInputData(any(Windmill.GlobalDataRequest.class))) .thenReturn( buildGlobalDataResponse(tag, ByteString.EMPTY, false, null), buildGlobalDataResponse(tag, ByteString.EMPTY, true, encodedIterable)); assertEquals( null, fetcher.fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier)); assertEquals( null, fetcher.fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier)); assertEquals( "data", fetcher .fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier) .orNull()); assertEquals( "data", fetcher .fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier) .orNull()); verify(server, times(2)).getSideInputData(buildGlobalDataRequest(tag, ByteString.EMPTY)); verifyNoMoreInteractions(server); }
@Test public void testFetchGlobalDataNull() throws Exception { StateFetcher fetcher = new StateFetcher(server); ByteString.Output stream = ByteString.newOutput(); ListCoder.of(VoidCoder.of()).encode(Arrays.asList((Void) null), stream, Coder.Context.OUTER); ByteString encodedIterable = stream.toByteString(); PCollectionView<Void> view = TestPipeline.create().apply(Create.empty(VoidCoder.of())).apply(View.asSingleton()); String tag = view.getTagInternal().getId(); when(server.getSideInputData(any(Windmill.GlobalDataRequest.class))) .thenReturn( buildGlobalDataResponse(tag, ByteString.EMPTY, false, null), buildGlobalDataResponse(tag, ByteString.EMPTY, true, encodedIterable)); assertEquals( null, fetcher.fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier)); assertEquals( null, fetcher.fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier)); assertEquals( null, fetcher .fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier) .orNull()); assertEquals( null, fetcher .fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.KNOWN_READY, readStateSupplier) .orNull()); verify(server, times(2)).getSideInputData(buildGlobalDataRequest(tag, ByteString.EMPTY)); verifyNoMoreInteractions(server); }
@Test public void testFetchGlobalDataCacheOverflow() throws Exception { Coder<List<String>> coder = ListCoder.of(StringUtf8Coder.of()); ByteString.Output stream = ByteString.newOutput(); coder.encode(Arrays.asList("data1"), stream, Coder.Context.OUTER); ByteString encodedIterable1 = stream.toByteString(); stream = ByteString.newOutput(); coder.encode(Arrays.asList("data2"), stream, Coder.Context.OUTER); ByteString encodedIterable2 = stream.toByteString(); Cache<StateFetcher.SideInputId, StateFetcher.SideInputCacheEntry> cache = CacheBuilder.newBuilder().build(); StateFetcher fetcher = new StateFetcher(server, cache); PCollectionView<String> view1 = TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton()); PCollectionView<String> view2 = TestPipeline.create().apply(Create.empty(StringUtf8Coder.of())).apply(View.asSingleton()); String tag1 = view1.getTagInternal().getId(); String tag2 = view2.getTagInternal().getId(); when(server.getSideInputData(any(Windmill.GlobalDataRequest.class))) .thenReturn( buildGlobalDataResponse(tag1, ByteString.EMPTY, true, encodedIterable1), buildGlobalDataResponse(tag2, ByteString.EMPTY, true, encodedIterable2), buildGlobalDataResponse(tag1, ByteString.EMPTY, true, encodedIterable1)); assertEquals( "data1", fetcher .fetchSideInput( view1, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier) .orNull()); assertEquals( "data2", fetcher .fetchSideInput( view2, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier) .orNull()); cache.invalidateAll(); assertEquals( "data1", fetcher .fetchSideInput( view1, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier) .orNull()); assertEquals( "data1", fetcher .fetchSideInput( view1, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier) .orNull()); ArgumentCaptor<Windmill.GlobalDataRequest> captor = ArgumentCaptor.forClass(Windmill.GlobalDataRequest.class); verify(server, times(3)).getSideInputData(captor.capture()); verifyNoMoreInteractions(server); assertThat( captor.getAllValues(), contains( buildGlobalDataRequest(tag1, ByteString.EMPTY), buildGlobalDataRequest(tag2, ByteString.EMPTY), buildGlobalDataRequest(tag1, ByteString.EMPTY))); }
@Test public void testEmptyFetchGlobalData() throws Exception { StateFetcher fetcher = new StateFetcher(server); ByteString encodedIterable = ByteString.EMPTY; PCollectionView<Long> view = TestPipeline.create() .apply(Create.empty(VarLongCoder.of())) .apply(Sum.longsGlobally().asSingletonView()); String tag = view.getTagInternal().getId(); when(server.getSideInputData(any(Windmill.GlobalDataRequest.class))) .thenReturn(buildGlobalDataResponse(tag, ByteString.EMPTY, true, encodedIterable)); assertEquals( 0L, (long) fetcher .fetchSideInput( view, GlobalWindow.INSTANCE, STATE_FAMILY, SideInputState.UNKNOWN, readStateSupplier) .orNull()); verify(server).getSideInputData(buildGlobalDataRequest(tag, ByteString.EMPTY)); verifyNoMoreInteractions(server); }