com.microsoft.azure.eventhubs.PartitionReceiver Java Examples

@OnStopped
public void tearDown() throws ProcessException {
    for (final PartitionReceiver receiver : partitionToReceiverMap.values()) {
        if (null != receiver) {
            receiver.close();
        }
    }

    partitionToReceiverMap.clear();
    try {
        if (null != eventHubClient) {
            eventHubClient.closeSync();
        }
    } catch (final ServiceBusException e) {
        throw new ProcessException(e);
    }
}
 

private static void consumeEvents(String ehName, String namespace, String keyName, String token)
    throws EventHubException, IOException, ExecutionException, InterruptedException {
  ConnectionStringBuilder connStr = new ConnectionStringBuilder().setNamespaceName(namespace)
      .setEventHubName(ehName)
      .setSasKeyName(keyName)
      .setSasKey(token);

  EventHubClient client = EventHubClient.createSync(connStr.toString(), Executors.newFixedThreadPool(10));

  EventHubRuntimeInformation runTimeInfo = client.getRuntimeInformation().get();
  int numPartitions = runTimeInfo.getPartitionCount();
  for (int partition = 0; partition < numPartitions; partition++) {
    PartitionReceiver receiver =
        client.createReceiverSync(EventHubClient.DEFAULT_CONSUMER_GROUP_NAME, String.valueOf(partition),
            EventPosition.fromStartOfStream());
    receiver.receive(10).handle((records, throwable) -> handleComplete(receiver, records, throwable));
  }
}
 

@OnStopped
public void tearDown() throws ProcessException {
    for (final PartitionReceiver receiver : partitionToReceiverMap.values()) {
        if (null != receiver) {
            receiver.close();
        }
    }

    partitionToReceiverMap.clear();
    try {
        if (null != eventHubClient) {
            eventHubClient.closeSync();
        }
        executor.shutdown();
    } catch (final EventHubException e) {
        throw new ProcessException(e);
    }
}
 

PartitionReceiver getReceiver(final ProcessContext context, final String partitionId) throws IOException, ServiceBusException, ExecutionException, InterruptedException {
    PartitionReceiver existingReceiver = partitionToReceiverMap.get(partitionId);
    if (existingReceiver != null) {
        return existingReceiver;
    }

    // we want to avoid allowing multiple threads to create Receivers simultaneously because that could result in
    // having multiple Receivers for the same partition. So if the map does not contain a receiver for this partition,
    // we will enter a synchronized block and check again (because once we enter the synchronized block, we know that no
    // other thread is creating a client). If within the synchronized block, we still do not have an entry in the map,
    // it is up to use to create the receiver, initialize it, and then put it into the map.
    // We do not use the putIfAbsent method in order to do a CAS operation here because we want to also initialize the
    // receiver if and only if it is not present in the map. As a result, we need to initialize the receiver and add it
    // to the map atomically. Hence, the synchronized block.
    synchronized (this) {
        existingReceiver = partitionToReceiverMap.get(partitionId);
        if (existingReceiver != null) {
            return existingReceiver;
        }

        final String consumerGroupName = context.getProperty(CONSUMER_GROUP).getValue();

        final PartitionReceiver receiver = eventHubClient.createReceiver(
                consumerGroupName,
                partitionId,
                configuredEnqueueTime == null ? Instant.now() : configuredEnqueueTime).get();

        receiver.setReceiveTimeout(receiverFetchTimeout == null ? Duration.ofMillis(60000) : receiverFetchTimeout);
        partitionToReceiverMap.put(partitionId, receiver);
        return receiver;

    }
}
 

@Override
protected PartitionReceiver getReceiver(final ProcessContext context, final String partitionId) throws IOException, ServiceBusException, ExecutionException, InterruptedException {
    if(getReceiverThrow){
        throw new IOException("Could not create receiver");
    }
    return null;
}
 

private static Object handleComplete(PartitionReceiver receiver, Iterable<EventData> records, Throwable throwable) {
  for (EventData record : records) {
    System.out.println(
        String.format("Partition %s, Event %s", receiver.getPartitionId(), new String(record.getBytes())));
  }

  receiver.receive(10).handle((r, t) -> handleComplete(receiver, r, t));
  return null;
}
 

private void renewPartitionReceiver(SystemStreamPartition ssp) {
  String streamId = config.getStreamId(ssp.getStream());
  EventHubClientManager eventHubClientManager = perPartitionEventHubManagers.get(ssp);
  String offset = streamPartitionOffsets.get(ssp);
  Integer partitionId = ssp.getPartition().getPartitionId();
  String consumerGroup = config.getStreamConsumerGroup(ssp.getSystem(), streamId);

  try {
    // Close current receiver
    streamPartitionReceivers.get(ssp).close().get(DEFAULT_SHUTDOWN_TIMEOUT_MILLIS, TimeUnit.MILLISECONDS);

    // Recreate receiver
    PartitionReceiver receiver = eventHubClientManager.getEventHubClient()
        .createReceiverSync(consumerGroup, partitionId.toString(),
            EventPosition.fromOffset(offset, !offset.equals(EventHubSystemConsumer.START_OF_STREAM)));

    receiver.setPrefetchCount(prefetchCount);

    // Timeout for EventHubClient receive
    receiver.setReceiveTimeout(DEFAULT_EVENTHUB_RECEIVER_TIMEOUT);

    // Create and start receiver thread with handler
    receiver.setReceiveHandler(streamPartitionHandlers.get(ssp));
    streamPartitionReceivers.put(ssp, receiver);
  } catch (Exception e) {
    eventHubNonTransientError.set(new SamzaException(
        String.format("Failed to recreate receiver for EventHubs after ReceiverHandlerError (ssp=%s)", ssp), e));
  }
}
 

@Test
public void testReceive() throws EventHubException {
  EventHubClientManagerFactory clientFactory = new EventHubClientManagerFactory();
  EventHubClientManager wrapper = clientFactory
          .getEventHubClientManager(SYSTEM_NAME, STREAM_NAME1, new EventHubConfig(createEventHubConfig()));
  wrapper.init();
  EventHubClient client = wrapper.getEventHubClient();
  PartitionReceiver receiver =
          client.createReceiverSync(EventHubClient.DEFAULT_CONSUMER_GROUP_NAME, "0",
                  EventPosition.fromStartOfStream());
  receiveMessages(receiver, 300);
}
 

private void receiveMessages(PartitionReceiver receiver, int numMessages) throws EventHubException {
  int count = 0;
  while (count < numMessages) {

    Iterable<EventData> messages = receiver.receiveSync(100);
    if (messages == null) {
      break;
    }
    for (EventData data : messages) {
      count++;
      LOG.info("Data" + new String(data.getBytes()));
    }
  }
}
 

@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);
}
 

PartitionReceiver getReceiver(final ProcessContext context, final String partitionId) throws IOException, EventHubException, ExecutionException, InterruptedException {
    PartitionReceiver existingReceiver = partitionToReceiverMap.get(partitionId);
    if (existingReceiver != null) {
        return existingReceiver;
    }

    // we want to avoid allowing multiple threads to create Receivers simultaneously because that could result in
    // having multiple Receivers for the same partition. So if the map does not contain a receiver for this partition,
    // we will enter a synchronized block and check again (because once we enter the synchronized block, we know that no
    // other thread is creating a client). If within the synchronized block, we still do not have an entry in the map,
    // it is up to use to create the receiver, initialize it, and then put it into the map.
    // We do not use the putIfAbsent method in order to do a CAS operation here because we want to also initialize the
    // receiver if and only if it is not present in the map. As a result, we need to initialize the receiver and add it
    // to the map atomically. Hence, the synchronized block.
    synchronized (this) {
        existingReceiver = partitionToReceiverMap.get(partitionId);
        if (existingReceiver != null) {
            return existingReceiver;
        }

        final String consumerGroupName = context.getProperty(CONSUMER_GROUP).getValue();

        final PartitionReceiver receiver = eventHubClient.createReceiver(
                consumerGroupName,
                partitionId,
                EventPosition.fromEnqueuedTime(
                        configuredEnqueueTime == null ? Instant.now() : configuredEnqueueTime)).get();

        receiver.setReceiveTimeout(receiverFetchTimeout == null ? Duration.ofMillis(60000) : receiverFetchTimeout);
        partitionToReceiverMap.put(partitionId, receiver);
        return receiver;

    }
}
 

@Override
protected PartitionReceiver getReceiver(final ProcessContext context, final String partitionId) throws IOException, EventHubException, ExecutionException, InterruptedException {
    if(getReceiverThrow){
        throw new IOException("Could not create receiver");
    }
    return null;
}
 

private synchronized void initializeEventHubsManagers() {
  LOG.info("Starting EventHubSystemConsumer. Count of SSPs registered: " + streamPartitionOffsets.entrySet().size());
  eventHubNonTransientError.set(null);
  // Create receivers for Event Hubs
  for (Map.Entry<SystemStreamPartition, String> entry : streamPartitionOffsets.entrySet()) {
    SystemStreamPartition ssp = entry.getKey();
    String streamId = config.getStreamId(ssp.getStream());
    Integer partitionId = ssp.getPartition().getPartitionId();
    String offset = entry.getValue();
    String consumerGroup = config.getStreamConsumerGroup(systemName, streamId);
    String namespace = config.getStreamNamespace(systemName, streamId);
    String entityPath = config.getStreamEntityPath(systemName, streamId);
    EventHubClientManager eventHubClientManager = createOrGetEventHubClientManagerForSSP(streamId, ssp);

    try {
      PartitionReceiver receiver;
      if (END_OF_STREAM.equals(offset)) {
        // If the offset is greater than the newest offset, use the use current Instant as
        // offset to fetch in Eventhub.
        receiver = eventHubClientManager.getEventHubClient()
            .createReceiver(consumerGroup, partitionId.toString(), EventPosition.fromEnqueuedTime(Instant.now())).get(DEFAULT_EVENTHUB_CREATE_RECEIVER_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
      } else {
        // EventHub will return the first message AFTER the offset that was specified in the fetch request.
        // If no such offset exists Eventhub will return an error.
        receiver = eventHubClientManager.getEventHubClient()
            .createReceiver(consumerGroup, partitionId.toString(),
                EventPosition.fromOffset(offset, /* inclusiveFlag */false)).get(DEFAULT_EVENTHUB_CREATE_RECEIVER_TIMEOUT.toMillis(), TimeUnit.MILLISECONDS);
      }

      receiver.setPrefetchCount(prefetchCount);

      PartitionReceiveHandler handler =
          new PartitionReceiverHandlerImpl(ssp, eventReadRates.get(streamId), eventByteReadRates.get(streamId),
              consumptionLagMs.get(streamId), readErrors.get(streamId), interceptors.getOrDefault(streamId, null),
              config.getMaxEventCountPerPoll(systemName));

      // Timeout for EventHubClient receive
      receiver.setReceiveTimeout(DEFAULT_EVENTHUB_RECEIVER_TIMEOUT);

      // Start the receiver thread
      receiver.setReceiveHandler(handler);

      streamPartitionHandlers.put(ssp, handler);
      streamPartitionReceivers.put(ssp, receiver);
    } catch (Exception e) {
      throw new SamzaException(
          String.format("Failed to create receiver for EventHubs: namespace=%s, entity=%s, partitionId=%d", namespace,
              entityPath, partitionId), e);
    }
    LOG.info(String.format("Connection successfully started for namespace=%s, entity=%s ", namespace, entityPath));
  }
}
 

/**
 * This method is here to try and isolate the Azure related code as the PartitionReceiver cannot be mocked
 * with PowerMock due to it being final. Unfortunately it extends a base class and does not implement an interface
 * so even if we create a MockPartitionReciver, it will not work as the two classes are orthogonal.
 *
 * @param context     - The processcontext for this processor
 * @param partitionId - The partition ID to retrieve a receiver by.
 * @return - Returns the events received from the EventBus.
 * @throws ProcessException -- If any exception is encountered, receiving events it is wrapped in a ProcessException
 *                          and then that exception is thrown.
 */
protected Iterable<EventData> receiveEvents(final ProcessContext context, final String partitionId) throws ProcessException {
    final PartitionReceiver receiver;
    try {
        receiver = getReceiver(context, partitionId);
        return receiver.receive(receiverFetchSize).get();
    } catch (final IOException | ServiceBusException | ExecutionException | InterruptedException e) {
        throw new ProcessException(e);
    }
}
 

/**
 * This method is here to try and isolate the Azure related code as the PartitionReceiver cannot be mocked
 * with PowerMock due to it being final. Unfortunately it extends a base class and does not implement an interface
 * so even if we create a MockPartitionReciver, it will not work as the two classes are orthogonal.
 *
 * @param context     - The processcontext for this processor
 * @param partitionId - The partition ID to retrieve a receiver by.
 * @return - Returns the events received from the EventBus.
 * @throws ProcessException -- If any exception is encountered, receiving events it is wrapped in a ProcessException
 *                          and then that exception is thrown.
 */
protected Iterable<EventData> receiveEvents(final ProcessContext context, final String partitionId) throws ProcessException {
    final PartitionReceiver receiver;
    try {
        receiver = getReceiver(context, partitionId);
        return receiver.receive(receiverFetchSize).get();
    } catch (final EventHubException | IOException | ExecutionException | InterruptedException e) {
        throw new ProcessException(e);
    }
}