KafkaSource extends AbstractPollableSource implements Configurable, BatchSizeSupported { @Override protected Status doProcess() throws EventDeliveryException { final String batchUUID = UUID.randomUUID().toString(); String kafkaKey; Event event; byte[] eventBody; try { final long nanoBatchStartTime = System.nanoTime(); final long batchStartTime = System.currentTimeMillis(); final long maxBatchEndTime = System.currentTimeMillis() + maxBatchDurationMillis; while (eventList.size() < batchUpperLimit && System.currentTimeMillis() < maxBatchEndTime) { if (it == null || !it.hasNext()) { long durMs = Math.max(0L, maxBatchEndTime - System.currentTimeMillis()); Duration duration = Duration.ofMillis(durMs); ConsumerRecords<String, byte[]> records = consumer.poll(duration); it = records.iterator(); if (rebalanceFlag.compareAndSet(true, false)) { break; } if (!it.hasNext()) { counter.incrementKafkaEmptyCount(); log.debug("Returning with backoff. No more data to read"); break; } } ConsumerRecord<String, byte[]> message = it.next(); kafkaKey = message.key(); if (useAvroEventFormat) { ByteArrayInputStream in = new ByteArrayInputStream(message.value()); decoder = DecoderFactory.get().directBinaryDecoder(in, decoder); if (!reader.isPresent()) { reader = Optional.of( new SpecificDatumReader<AvroFlumeEvent>(AvroFlumeEvent.class)); } AvroFlumeEvent avroevent = reader.get().read(null, decoder); eventBody = avroevent.getBody().array(); headers = toStringMap(avroevent.getHeaders()); } else { eventBody = message.value(); headers.clear(); headers = new HashMap<String, String>(4); } if (!headers.containsKey(KafkaSourceConstants.TIMESTAMP_HEADER)) { headers.put(KafkaSourceConstants.TIMESTAMP_HEADER, String.valueOf(System.currentTimeMillis())); } if (setTopicHeader && !headers.containsKey(topicHeader)) { headers.put(topicHeader, message.topic()); } if (!headers.containsKey(KafkaSourceConstants.PARTITION_HEADER)) { headers.put(KafkaSourceConstants.PARTITION_HEADER, String.valueOf(message.partition())); } if (!headers.containsKey(OFFSET_HEADER)) { headers.put(OFFSET_HEADER, String.valueOf(message.offset())); } if (kafkaKey != null) { headers.put(KafkaSourceConstants.KEY_HEADER, kafkaKey); } if (log.isTraceEnabled()) { if (LogPrivacyUtil.allowLogRawData()) { log.trace("Topic: {} Partition: {} Message: {}", new String[]{ message.topic(), String.valueOf(message.partition()), new String(eventBody) }); } else { log.trace("Topic: {} Partition: {} Message arrived.", message.topic(), String.valueOf(message.partition())); } } event = EventBuilder.withBody(eventBody, headers); eventList.add(event); if (log.isDebugEnabled()) { log.debug("Waited: {} ", System.currentTimeMillis() - batchStartTime); log.debug("Event #: {}", eventList.size()); } tpAndOffsetMetadata.put(new TopicPartition(message.topic(), message.partition()), new OffsetAndMetadata(message.offset() + 1, batchUUID)); } if (eventList.size() > 0) { counter.addToKafkaEventGetTimer((System.nanoTime() - nanoBatchStartTime) / (1000 * 1000)); counter.addToEventReceivedCount((long) eventList.size()); getChannelProcessor().processEventBatch(eventList); counter.addToEventAcceptedCount(eventList.size()); if (log.isDebugEnabled()) { log.debug("Wrote {} events to channel", eventList.size()); } eventList.clear(); if (!tpAndOffsetMetadata.isEmpty()) { long commitStartTime = System.nanoTime(); consumer.commitSync(tpAndOffsetMetadata); long commitEndTime = System.nanoTime(); counter.addToKafkaCommitTimer((commitEndTime - commitStartTime) / (1000 * 1000)); tpAndOffsetMetadata.clear(); } return Status.READY; } return Status.BACKOFF; } catch (Exception e) { log.error("KafkaSource EXCEPTION, {}", e); counter.incrementEventReadOrChannelFail(e); return Status.BACKOFF; } }  @Override long getBatchSize();  }
@Test public void testErrorCounters() throws InterruptedException, EventDeliveryException { context.put(TOPICS, topic0); context.put(BATCH_SIZE, "1"); kafkaSource.configure(context); ChannelProcessor cp = Mockito.mock(ChannelProcessor.class); doThrow(new ChannelException("dummy")).doThrow(new RuntimeException("dummy")) .when(cp).processEventBatch(any(List.class)); kafkaSource.setChannelProcessor(cp); startKafkaSource(); Thread.sleep(500L); kafkaServer.produce(topic0, "", "hello, world"); Thread.sleep(500L); kafkaSource.doProcess(); kafkaSource.doProcess(); SourceCounter sc = (SourceCounter) Whitebox.getInternalState(kafkaSource, "counter"); Assert.assertEquals(1, sc.getChannelWriteFail()); Assert.assertEquals(1, sc.getEventReadFail()); kafkaSource.stop(); }