Source:
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) { Map<String, SystemStreamMetadata> requestedMetadata = new HashMap<>(); Map<String, CompletableFuture<EventHubRuntimeInformation>> ehRuntimeInfos = new HashMap<>(); streamNames.forEach((streamName) -> { if (!streamPartitions.containsKey(streamName)) { EventHubClientWrapper eventHubClientWrapper = getStreamEventHubClient(streamName); ehRuntimeInfos.put(streamName, eventHubClientWrapper.getEventHubClient().getRuntimeInformation()); } }); ehRuntimeInfos.forEach((streamName, ehRuntimeInfo) -> { if (!streamPartitions.containsKey(streamName)) { try { EventHubRuntimeInformation ehInfo = ehRuntimeInfo.get(eventHubConfig.getRuntimeInfoWaitTimeMS(), TimeUnit.MILLISECONDS); streamPartitions.put(streamName, ehInfo.getPartitionIds()); } catch (InterruptedException | ExecutionException e) { String msg = String.format("Error while fetching EventHubRuntimeInfo for System:%s, Stream:%s", systemName, streamName); LOG.error(msg, e); throw new SamzaException(msg); } catch (TimeoutException e) { String msg = String.format("Timed out while fetching EventHubRuntimeInfo for System:%s, Stream:%s", systemName, streamName); LOG.error(msg, e); throw new SamzaException(msg); } } requestedMetadata.put(streamName, new SystemStreamMetadata(streamName, getPartitionMetadata(streamName, streamPartitions.get(streamName)))); }); return requestedMetadata; }

Target:
public Map<String, SystemStreamMetadata> getSystemStreamMetadata(Set<String> streamNames) { final Map<String, SystemStreamMetadata> requestedMetadata = new HashMap<>(); final Map<String, CompletableFuture<EventHubRuntimeInformation>> ehRuntimeInfos = new HashMap<>(); streamNames.forEach((streamName) -> { if (!streamPartitions.containsKey(streamName)) { EventHubClientManager eventHubClientManager = getOrCreateStreamEventHubClient(streamName); CompletableFuture<EventHubRuntimeInformation> runtimeInfo = eventHubClientManager.getEventHubClient() .getRuntimeInformation(); ehRuntimeInfos.put(streamName, runtimeInfo); } }); try { ehRuntimeInfos.forEach((streamName, ehRuntimeInfo) -> { if (!streamPartitions.containsKey(streamName)) { LOG.debug(String.format("Partition ids for Stream=%s not found", streamName)); try { long timeoutMs = eventHubConfig.getRuntimeInfoWaitTimeMS(systemName); EventHubRuntimeInformation ehInfo = ehRuntimeInfo.get(timeoutMs, TimeUnit.MILLISECONDS); LOG.debug(String.format("Adding partition ids for Stream=%s", streamName)); streamPartitions.put(streamName, ehInfo.getPartitionIds()); } catch (InterruptedException | ExecutionException | TimeoutException e) { String msg = String.format("Error while fetching EventHubRuntimeInfo for System:%s, Stream:%s", systemName, streamName); throw new SamzaException(msg); } } String[] partitionIds = streamPartitions.get(streamName); Map<Partition, SystemStreamPartitionMetadata> sspMetadataMap = getPartitionMetadata(streamName, partitionIds); SystemStreamMetadata systemStreamMetadata = new SystemStreamMetadata(streamName, sspMetadataMap); requestedMetadata.put(streamName, systemStreamMetadata); }); } finally { eventHubClients.forEach((streamName, client) -> client.close(DEFAULT_SHUTDOWN_TIMEOUT_MILLIS)); eventHubClients.clear(); } return requestedMetadata; }