nstream.adapter.kafka.KafkaIngestingAgent Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of nstream-adapter-kafka Show documentation
Templates for consuming from and producing to Kafka topics with Swim
There is a newer version: 4.15.23
Show newest version
// Copyright 2015-2024 Nstream, inc.
//
// Licensed under the Redis Source Available License 2.0 (RSALv2) Agreement;
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://redis.com/legal/rsalv2-agreement/
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package nstream.adapter.kafka;

import java.time.Duration;
import java.util.Collection;
import java.util.Set;
import nstream.adapter.common.ext.KafkaIngressSettings;
import nstream.adapter.common.ingress.IngestorMetricsAgent;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.common.TopicPartition;
import swim.concurrent.AbstractTask;
import swim.concurrent.TaskRef;
import swim.concurrent.TimerRef;
import swim.structure.Value;

public abstract class KafkaIngestingAgent
    extends IngestorMetricsAgent> {

  protected static final long DISCOVERY_POLL_MS = 5000L;
  protected static final int DISCOVERY_POLL_ATTEMPTS = 6;

  protected volatile Consumer kafkaConsumer;
  protected volatile TimerRef pollTimer;
  protected volatile TaskRef discoveryTask;
  protected volatile TaskRef pollTask;

  public KafkaIngestingAgent() {
  }

  protected TimerRef pollTimer() {
    return this.pollTimer;
  }

  protected TaskRef pollTask() {
    return this.pollTask;
  }

  protected Consumer kafkaConsumer() {
    return this.kafkaConsumer;
  }

  protected void prepareConsumer() {
    cancel();
    loadSettings("kafkaIngressConf");
    // each of the next two lines is blocking and exception-prone
    this.kafkaConsumer = KafkaAdapterUtils.createConsumer(this.ingressSettings);
    // FIXME: support assign() as a configurable option from KafkaIngressSettings
    this.kafkaConsumer.subscribe(this.ingressSettings.topics(), new ConsumerRebalanceListener() {

      @Override
      public void onPartitionsRevoked(Collection partitions) {
        info(nodeUri() + ": revoked partitions " + partitions);
      }

      @Override
      public void onPartitionsAssigned(Collection partitions) {
        info(nodeUri() + ": assigned partitions " + partitions);
      }

    });
  }

  protected void releaseConsumer() {
    if (this.kafkaConsumer != null) {
      try {
        this.kafkaConsumer.close();
      } catch (Exception e) {
        didFail(new RuntimeException(nodeUri() + ": exception observed in releasing kafkaConsumer "
            + "(resources were probably cleared regardless)", e));
      }
      this.kafkaConsumer = null;
    }
  }

  protected void simpleReceptionStage() {
    this.pollTask = prepareLoop(this::pollTask, () -> {
      final ConsumerRecords batch = poll();
      ingestBatch(batch);
    });
    this.pollTask.cue();
  }

  protected void discoveryReceptionStage() {
    this.discoveryTask = asyncStage().task(new AbstractTask() {

      private volatile int attempts = 0;

      @Override
      public void runTask() {
        if (this.attempts++ < DISCOVERY_POLL_ATTEMPTS) {
          final ConsumerRecords batch = discoveryPoll();
          final Set assignment = KafkaIngestingAgent.this.kafkaConsumer.assignment();
          if (!assignment.isEmpty()) {
            info(nodeUri() + ": discovery poll is aware of kafkaConsumer assignment " + assignment);
            if (!batch.isEmpty()) {
              ingestBatch(batch);
            }
            simpleReceptionStage();
            try {
              KafkaIngestingAgent.this.discoveryTask.cancel();
              KafkaIngestingAgent.this.discoveryTask = null;
            } catch (Exception e) {
              didFail(e);
            }
            return;
          }
        } else {
          warn(nodeUri() + ": kafkaConsumer received no assignment after " + DISCOVERY_POLL_ATTEMPTS * DISCOVERY_POLL_MS
              + " ms, possibly indicating invalid bootstrap servers or too many consumers in group. Retrying...");
          this.attempts = 0;
        }
        KafkaIngestingAgent.this.discoveryTask.cue();
      }

      @Override
      public boolean taskWillBlock() {
        return true;
      }

    });
    this.discoveryTask.cue();
  }

  protected ConsumerRecords discoveryPoll() {
    return this.kafkaConsumer.poll(Duration.ofMillis(DISCOVERY_POLL_MS));
  }

  protected ConsumerRecords poll() {
    return this.kafkaConsumer.poll(Duration.ofMillis(this.ingressSettings.pollTimeoutMillis()));
  }

  protected void ingestBatch(ConsumerRecords batch) {
    for (ConsumerRecord record : batch) {
      ingestOrCancel(record);
    }
  }

  @Override
  protected void cancel() {
    if (this.pollTask != null) {
      this.pollTask.cancel();
      this.pollTask = null;
    }
    if (this.pollTimer != null) {
      this.pollTimer.cancel();
      this.pollTimer = null;
    }
    releaseConsumer();
  }

  @Override
  protected KafkaIngressSettings parseIngressSettings(Value prop) {
    final KafkaIngressSettings settings = KafkaIngressSettings.form().cast(prop);
    return settings == null ? KafkaIngressSettings.defaultSettings() : settings;
  }

  @Override
  protected void stageReception() {
    prepareConsumer();
    // This line can be replaced with simpleReceptionStage(), but you lose
    // lifecycle logging upon nonempty kafkaConsumer.assignment(), AND you risk
    // silent failure if ingressSettings.pollTimeoutMillis() is too short or if
    // the bootstrap servers are DNS-resolvable but invalid.
    discoveryReceptionStage();
  }

}