PersistedMessageCallback.java

/*
 * Copyright © 2014 Cask Data, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License"); you may not
 * use this file except in compliance with the License. You may obtain a copy of
 * the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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 co.cask.cdap.metrics.process;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Maps;
import org.apache.twill.kafka.client.FetchedMessage;
import org.apache.twill.kafka.client.KafkaConsumer;
import org.apache.twill.kafka.client.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * A {@link KafkaConsumer.MessageCallback} that persists offset information into a VCTable while
 * delegating the actual message consumption to another {@link KafkaConsumer.MessageCallback}.
 */
public final class PersistedMessageCallback implements KafkaConsumer.MessageCallback {

  private static final Logger LOG = LoggerFactory.getLogger(PersistedMessageCallback.class);

  private final KafkaConsumer.MessageCallback delegate;
  private final KafkaConsumerMetaTable metaTable;
  private final int persistThreshold;
  private final Map<TopicPartition, Long> offsets;
  private final AtomicInteger messageCount;

  /**
   * Constructs a {@link PersistedMessageCallback} which delegates to the given callback for actual action while
   * persisting offsets information in to the given meta table when number of messages has been processed based
   * on the persistThreshold.
   */
  public PersistedMessageCallback(KafkaConsumer.MessageCallback delegate,
                                  KafkaConsumerMetaTable metaTable,
                                  int persistThreshold) {
    this.delegate = delegate;
    this.metaTable = metaTable;
    this.persistThreshold = persistThreshold;
    this.offsets = Maps.newConcurrentMap();
    this.messageCount = new AtomicInteger();
  }

  @Override
  public void onReceived(Iterator<FetchedMessage> messages) {
    delegate.onReceived(new OffsetTrackingIterator(messages));
    if (messageCount.get() >= persistThreshold) {
      messageCount.set(0);
      persistOffsets();
    }
  }

  @Override
  public void finished() {
    try {
      delegate.finished();
    } finally {
      // Save the offset
      persistOffsets();
    }
  }

  private void persistOffsets() {
    try {
      metaTable.save(ImmutableMap.copyOf(offsets));
    } catch (Exception e) {
      // Simple log and ignore the error.
      LOG.error("Failed to persist consumed message offset. {}", e.getMessage(), e);
    }
  }

  /**
   * Inner help class to track offsets of {@link FetchedMessage} being consumed and
   * persist to meta table when persistThreshold is reached.
   */
  private final class OffsetTrackingIterator implements Iterator<FetchedMessage> {

    private final Iterator<FetchedMessage> delegate;
    private TopicPartition lastTopicPartition;
    private long lastOffset = -1;

    OffsetTrackingIterator(Iterator<FetchedMessage> delegate) {
      this.delegate = delegate;
    }

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

    @Override
    public FetchedMessage next() {
      FetchedMessage message = delegate.next();
      lastTopicPartition = message.getTopicPartition();
      lastOffset = message.getNextOffset();
      messageCount.incrementAndGet();
      recordOffset();
      return message;
    }

    @Override
    public void remove() {
      delegate.remove();
    }

    private void recordOffset() {
      if (lastOffset >= 0) {
        offsets.put(lastTopicPartition, lastOffset);
      }
    }
  }
}