com.noleme.flow.connect.kafka.generator.KafkaConsumerGenerator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of noleme-flow-connect-kafka Show documentation
The newest version!
package com.noleme.flow.connect.kafka.generator;

import com.noleme.flow.actor.generator.Generator;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Duration;
import java.util.Collection;
import java.util.LinkedList;
import java.util.List;
import java.util.Queue;

/**
 * @author Pierre Lecerf ([email protected])
 */
public class KafkaConsumerGenerator implements Generator
{
    private static final Logger logger = LoggerFactory.getLogger(KafkaConsumerGenerator.class);

    private final Consumer consumer;
    private final Duration pollTimeout;
    private final Collection topics;
    private final Queue> queue;
    private boolean isActive;

    public KafkaConsumerGenerator(Consumer consumer, String... topics)
    {
        this(consumer, Duration.ofMillis(500), topics);
    }

    public KafkaConsumerGenerator(Consumer consumer, Duration pollTimeout, String... topics)
    {
        this(consumer, pollTimeout, List.of(topics));
    }

    public KafkaConsumerGenerator(Consumer consumer, Duration pollTimeout, Collection topics)
    {
        this.topics = topics;
        this.consumer = consumer;
        this.pollTimeout = pollTimeout;
        this.queue = new LinkedList<>();
        this.isActive = true;

        this.consumer.subscribe(topics);
    }

    @Override
    public boolean hasNext()
    {
        return this.isActive;
    }

    @Override
    public V generate()
    {
        if (!this.queue.isEmpty())
        {
            logger.debug("Polling message from in-memory queue (queue size: {})", this.queue.size());
            return this.poll();
        }

        while (this.isActive && this.queue.isEmpty())
        {
            for (ConsumerRecord record : this.consumer.poll(this.pollTimeout))
                this.queue.add(record);
        }

        return this.poll();
    }

    private V poll()
    {
        logger.info("Receiving message from kafka topics: {}", this.topics);
        var record = this.queue.poll();
        return record != null ? record.value() : null; //This should never be null anyway due to checks in generate()
    }

    synchronized public void deactivate()
    {
        this.isActive = false;
    }
}