Examples with EventHubClientManager org.apache.samza.system.eventhub.EventHubClientManager used on opensource projects

Example 1 with EventHubClientManager

use of org.apache.samza.system.eventhub.EventHubClientManager in project samza by apache.

the class TestEventHubSystemAdmin method testStartpointResolverShouldResolveTheStartpointTimestampToCorrectOffset.

@Test
public void testStartpointResolverShouldResolveTheStartpointTimestampToCorrectOffset() throws EventHubException {
    // Initialize variables required for testing.
    EventHubSystemAdmin mockEventHubSystemAdmin = Mockito.mock(EventHubSystemAdmin.class);
    EventHubConfig eventHubConfig = Mockito.mock(EventHubConfig.class);
    SystemStreamPartition systemStreamPartition = new SystemStreamPartition("test-system", "test-stream", new Partition(0));
    String mockedOffsetToReturn = "100";
    // Setup the mock variables.
    EventHubClientManager mockEventHubClientManager = Mockito.mock(EventHubClientManager.class);
    EventHubClient mockEventHubClient = Mockito.mock(EventHubClient.class);
    PartitionReceiver mockPartitionReceiver = Mockito.mock(PartitionReceiver.class);
    EventData mockEventData = Mockito.mock(EventData.class);
    EventData.SystemProperties mockSystemProperties = Mockito.mock(EventData.SystemProperties.class);
    // Configure the mock variables to return the appropriate values.
    Mockito.when(mockEventHubSystemAdmin.getOrCreateStreamEventHubClient("test-stream")).thenReturn(mockEventHubClientManager);
    Mockito.when(mockEventHubClientManager.getEventHubClient()).thenReturn(mockEventHubClient);
    Mockito.when(mockEventHubClient.createReceiverSync(Mockito.anyString(), Mockito.anyString(), Mockito.any())).thenReturn(mockPartitionReceiver);
    Mockito.when(mockPartitionReceiver.receiveSync(1)).thenReturn(Arrays.asList(mockEventData));
    Mockito.when(mockEventData.getSystemProperties()).thenReturn(mockSystemProperties);
    Mockito.when(mockSystemProperties.getOffset()).thenReturn(mockedOffsetToReturn);
    // Test the Offset resolver.
    EventHubSamzaOffsetResolver resolver = new EventHubSamzaOffsetResolver(mockEventHubSystemAdmin, eventHubConfig);
    String resolvedOffset = resolver.visit(systemStreamPartition, new StartpointTimestamp(100L));
    Assert.assertEquals(mockedOffsetToReturn, resolvedOffset);
}

Example 2 with EventHubClientManager

use of org.apache.samza.system.eventhub.EventHubClientManager in project samza by apache.

the class EventHubSystemProducer method createOrGetEventHubClientManagerForPartition.

private EventHubClientManager createOrGetEventHubClientManagerForPartition(String streamId, int partitionId) {
    Map<Integer, EventHubClientManager> perStreamMap = perPartitionEventHubClients.computeIfAbsent(streamId, key -> new HashMap<>());
    EventHubClientManager eventHubClientManager;
    if (config.getPerPartitionConnection(systemName)) {
        // will create one EventHub client per partition
        if (perStreamMap.containsKey(partitionId)) {
            LOG.warn(String.format("Trying to create new EventHubClientManager for partition=%d. But one already exists", partitionId));
            eventHubClientManager = perStreamMap.get(partitionId);
        } else {
            LOG.info(String.format("Creating EventHub client manager for streamId=%s, partitionId=%d: ", streamId, partitionId));
            eventHubClientManager = eventHubClientManagerFactory.getEventHubClientManager(systemName, streamId, config);
            eventHubClientManager.init();
            perStreamMap.put(partitionId, eventHubClientManager);
        }
    } else {
        // will share one EventHub client per stream
        eventHubClientManager = perStreamEventHubClientManagers.get(streamId);
        perStreamMap.put(partitionId, eventHubClientManager);
    }
    Validate.notNull(eventHubClientManager, String.format("Fail to create or get EventHubClientManager for streamId=%s, partitionId=%d", streamId, partitionId));
    return eventHubClientManager;
}

Example 3 with EventHubClientManager

use of org.apache.samza.system.eventhub.EventHubClientManager in project samza by apache.

the class EventHubSystemProducer method init.

private void init() {
    LOG.info("Initializing EventHubSystemProducer");
    // Fetches the stream ids
    List<String> streamIds = config.getStreams(systemName);
    // partition count)
    for (String streamId : streamIds) {
        EventHubClientManager ehClient = eventHubClientManagerFactory.getEventHubClientManager(systemName, streamId, config);
        perStreamEventHubClientManagers.put(streamId, ehClient);
        ehClient.init();
    }
    // Create partition senders if required
    if (PartitioningMethod.PARTITION_KEY_AS_PARTITION.equals(partitioningMethod)) {
        // Create all partition senders
        perStreamEventHubClientManagers.forEach((streamId, samzaEventHubClient) -> {
            EventHubClient ehClient = samzaEventHubClient.getEventHubClient();
            try {
                Map<Integer, PartitionSender> partitionSenders = new HashMap<>();
                long timeoutMs = config.getRuntimeInfoWaitTimeMS(systemName);
                Integer numPartitions = ehClient.getRuntimeInformation().get(timeoutMs, TimeUnit.MILLISECONDS).getPartitionCount();
                for (int i = 0; i < numPartitions; i++) {
                    String partitionId = String.valueOf(i);
                    EventHubClientManager perPartitionClientManager = createOrGetEventHubClientManagerForPartition(streamId, i);
                    PartitionSender partitionSender = perPartitionClientManager.getEventHubClient().createPartitionSender(partitionId).get(DEFAULT_CREATE_PARTITION_SENDER_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);
                    partitionSenders.put(i, partitionSender);
                }
                streamPartitionSenders.put(streamId, partitionSenders);
            } catch (InterruptedException | ExecutionException | TimeoutException e) {
                String msg = "Failed to fetch number of Event Hub partitions for partition sender creation";
                throw new SamzaException(msg, e);
            } catch (EventHubException | IllegalArgumentException e) {
                String msg = "Creation of partition sender failed with exception";
                throw new SamzaException(msg, e);
            }
        });
    }
    isInitialized = true;
    LOG.info("EventHubSystemProducer initialized.");
}

Example 4 with EventHubClientManager

use of org.apache.samza.system.eventhub.EventHubClientManager in project samza by apache.

the class EventHubSystemAdmin method getOrCreateStreamEventHubClient.

EventHubClientManager getOrCreateStreamEventHubClient(String streamName) {
    if (!eventHubClients.containsKey(streamName)) {
        LOG.info(String.format("Creating EventHubClient for Stream=%s", streamName));
        EventHubClientManager eventHubClientManager = eventHubClientManagerFactory.getEventHubClientManager(systemName, streamName, eventHubConfig);
        eventHubClientManager.init();
        eventHubClients.put(streamName, eventHubClientManager);
    }
    return eventHubClients.get(streamName);
}

Example 5 with EventHubClientManager

use of org.apache.samza.system.eventhub.EventHubClientManager in project samza by apache.

the class EventHubSystemAdmin method getPartitionMetadata.

private Map<Partition, SystemStreamPartitionMetadata> getPartitionMetadata(String streamName, String[] partitionIds) {
    EventHubClientManager eventHubClientManager = getOrCreateStreamEventHubClient(streamName);
    Map<Partition, SystemStreamPartitionMetadata> sspMetadataMap = new HashMap<>();
    List<CompletableFuture<PartitionRuntimeInformation>> futureList = new ArrayList<>();
    for (String partition : partitionIds) {
        CompletableFuture<PartitionRuntimeInformation> partitionRuntimeInfo = eventHubClientManager.getEventHubClient().getPartitionRuntimeInformation(partition);
        futureList.add(partitionRuntimeInfo);
        partitionRuntimeInfo.thenAccept(ehPartitionInfo -> {
            LOG.info(printPartitionRuntimeInfo(ehPartitionInfo));
            // Set offsets
            String startingOffset = EventHubSystemConsumer.START_OF_STREAM;
            String newestOffset = ehPartitionInfo.getLastEnqueuedOffset();
            String upcomingOffset = EventHubSystemConsumer.END_OF_STREAM;
            SystemStreamPartitionMetadata sspMetadata = new SystemStreamPartitionMetadata(startingOffset, newestOffset, upcomingOffset);
            sspMetadataMap.put(new Partition(Integer.parseInt(partition)), sspMetadata);
        });
    }
    CompletableFuture<Void> futureGroup = CompletableFuture.allOf(futureList.toArray(new CompletableFuture[futureList.size()]));
    long timeoutMs = eventHubConfig.getRuntimeInfoWaitTimeMS(systemName);
    try {
        futureGroup.get(timeoutMs, TimeUnit.MILLISECONDS);
    } catch (InterruptedException | ExecutionException | TimeoutException e) {
        String msg = String.format("Error while fetching EventHubPartitionRuntimeInfo for System:%s, Stream:%s", systemName, streamName);
        LOG.error(msg, e);
        throw new SamzaException(msg, e);
    }
    return sspMetadataMap;
}