PartitionReader implements Comparable<PartitionReader> { public List<JetstreamEvent> readEvents() throws OffsetOutOfRangeException { List<JetstreamEvent> events = new ArrayList<JetstreamEvent>(); if (m_nextBatchSizeBytes < 0) m_nextBatchSizeBytes = m_config.getBatchSizeBytes(); if (m_nextBatchSizeBytes == 0) { m_nextBatchSizeBytes = m_config.getBatchSizeBytes(); return events; } for (int i = 0; i < 2; i++) { for (int j = 0; j < m_retries; j++) { FetchRequest req = new FetchRequestBuilder() .clientId(m_clientId) .addFetch(m_topic, m_partition, m_readOffset, m_nextBatchSizeBytes).build(); FetchResponse fetchResponse = null; try { fetchResponse = m_consumer.fetch(req); } catch (Exception e) { LOGGER.error( "Error occurs when fetching.", e); continue; } if (fetchResponse == null) continue; if (fetchResponse.hasError()) { short code = fetchResponse.errorCode(m_topic, m_partition); if (code == ErrorMapping.OffsetOutOfRangeCode()) { long smallest = -1L, largest = -1L; try { smallest = fetchResetOffset(SmallestTimeString()); largest = fetchResetOffset(LargestTimeString()); } catch (Exception e) { } throw new OffsetOutOfRangeException(m_clientId + ": readOffset=" + m_readOffset + " smallest=" + smallest + " largest=" + largest); } else continue; } else { ByteBufferMessageSet messageSet = fetchResponse.messageSet( m_topic, m_partition); boolean hasMessage = messageSet.iterator().hasNext(); if (!hasMessage && !readToTheEnd()) { m_nextBatchSizeBytes = Math.min( m_nextBatchSizeBytes * 2, m_config.getMaxBatchSizeBytes()); continue; } for (MessageAndOffset messageAndOffset : messageSet) { long currentOffset = messageAndOffset.offset(); if (currentOffset < m_readOffset) { continue; } m_readOffset = messageAndOffset.nextOffset(); Message message = messageAndOffset.message(); ByteBuffer k = message.key(); ByteBuffer p = message.payload(); byte[] key = null; if (k != null) { key = new byte[k.limit()]; k.get(key); } byte[] payload = new byte[p.limit()]; p.get(payload); JetstreamEvent event = m_serializer .decode(key, payload); EventKafkaMetadata meta = new EventKafkaMetadata( m_topic, m_partition, currentOffset); event.put(kmKey, meta.encode()); events.add(event); } m_nextBatchSizeBytes = m_config.getBatchSizeBytes(); return events; } } reinit(); } throw new RuntimeException("Fail to read events from " + m_clientId + " with offset " + m_readOffset); } PartitionReader(String topic, int partition,
			KafkaConsumerConfig config,
			KafkaController.ZkConnector zkConnector,
			KafkaMessageSerializer serializer); void initOffset(boolean isNew); void resetOffset(); void revertOffset(); long getReadOffset(); void setReadOffset(long offset); void setNextBatchSizeBytes(int nextBatchSizeBytes); boolean isTaken(); void setToRelease(); void setLost(); boolean isLost(); String getTopic(); int getPartition(); String getId(); long calcWaitTime(); long getNextFetchInMs(); void setNextFetchInMs(long nextFetchInMs); void fistIdle(); long getIdleTime(); boolean isIdled(); void setOnIdled(); boolean isOnIdled(); void clearIdleStats(); List<JetstreamEvent> readEvents(); void commitOffset(); void close(); long readZkOffset(); @Override int compareTo(PartitionReader p);  }
@Test public void testReadOffsetOutOfRange() throws Exception { long smallest = 100, largest = 10000; SimpleConsumer mockConsumer = mock(SimpleConsumer.class); mockSimpleConsumerForEarliestOffset(mockConsumer, topic, partition, smallest); mockSimpleConsumerForLatestOffset(mockConsumer, topic, partition, largest); mockSimpleConsumerForOffsetOutOfRange(mockConsumer, topic, partition); ReflectFieldUtil.setField(reader, "m_consumer", mockConsumer); boolean thrown = false; try { reader.readEvents(); } catch (Exception e) { assertTrue(e instanceof OffsetOutOfRangeException); thrown = true; } assertTrue(thrown); }
@Test public void testFetchException() throws Exception { SimpleConsumer mockConsumer = mock(SimpleConsumer.class); mockSimpleConsumerForFetchException(mockConsumer); ReflectFieldUtil.setField(reader, "m_consumer", mockConsumer); boolean thrown = false; try { reader.readEvents(); } catch (Exception e) { assertTrue(e instanceof RuntimeException); thrown = true; } assertTrue(thrown); }
@Test public void testReadEventsAndFetchNull() throws Exception { SimpleConsumer mockConsumer = mock(SimpleConsumer.class); this.mockSimpleConsumerForNull(mockConsumer); ReflectFieldUtil.setField(reader, "m_consumer", mockConsumer); try { List<JetstreamEvent> events = reader.readEvents(); } catch (Exception e) { } verify(mockConsumer, atLeast(2)).fetch(argThat(new IsFetchRequest())); }
@Test public void testReadEvents() throws Exception { int expectedCount = (int) (consumerConfig.getBatchSizeBytes() / eachEventInBytes); long readOffset = 1250; reader.setReadOffset(readOffset); SimpleConsumer mockConsumer = mock(SimpleConsumer.class); mockSimpleConsumerForRead(mockConsumer, topic, partition, readOffset, consumerConfig.getBatchSizeBytes()); ReflectFieldUtil.setField(reader, "m_consumer", mockConsumer); reader.setNextBatchSizeBytes(-100); List<JetstreamEvent> events = reader.readEvents(); assertEquals(expectedCount, events.size()); assertEquals(readOffset + expectedCount, reader.getReadOffset()); }