com.wavefront.agent.queueing.FileBasedTaskQueue Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of proxy Show documentation
Service for batching and relaying metric traffic to Wavefront
There is a newer version: 4.36
package com.wavefront.agent.queueing;

import com.wavefront.agent.data.DataSubmissionTask;
import com.wavefront.common.Utils;
import java.io.IOException;
import java.util.Iterator;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Logger;
import javax.annotation.Nonnull;
import javax.annotation.Nullable;

/**
 * Implements proxy-specific {@link TaskQueue} interface as a wrapper over {@link QueueFile}.
 *
 * @param  type of objects stored.
 * @author [email protected]
 */
public class FileBasedTaskQueue> implements TaskQueue {
  private static final Logger log = Logger.getLogger(FileBasedTaskQueue.class.getCanonicalName());

  private final DirectByteArrayOutputStream bytes = new DirectByteArrayOutputStream();

  private volatile T head;

  private final AtomicLong currentWeight = new AtomicLong();
  private final QueueFile queueFile;
  private final TaskConverter taskConverter;

  /**
   * @param queueFile file backing the queue
   * @param taskConverter task converter
   */
  public FileBasedTaskQueue(QueueFile queueFile, TaskConverter taskConverter) {
    this.queueFile = queueFile;
    this.taskConverter = taskConverter;
    log.fine("Enumerating queue");
    this.queueFile
        .iterator()
        .forEachRemaining(
            task -> {
              Integer weight = taskConverter.getWeight(task);
              if (weight != null) {
                currentWeight.addAndGet(weight);
              }
            });
    log.fine("Enumerated: " + currentWeight.get() + " items in " + queueFile.size() + " tasks");
  }

  @Override
  public T peek() {
    try {
      if (this.head != null) {
        return this.head;
      }
      byte[] task = queueFile.peek();
      if (task == null) return null;
      this.head = taskConverter.fromBytes(task);
      return this.head;
    } catch (IOException ex) {
      throw Utils.throwAny(ex);
    }
  }

  @Override
  public void add(@Nonnull T entry) throws IOException {
    bytes.reset();
    taskConverter.serializeToStream(entry, bytes);
    queueFile.add(bytes.getArray(), 0, bytes.size());
    currentWeight.addAndGet(entry.weight());
  }

  @Override
  public void clear() throws IOException {
    queueFile.clear();
    this.head = null;
    this.currentWeight.set(0);
  }

  @Override
  public void remove() throws IOException {
    if (this.head == null) {
      byte[] task = queueFile.peek();
      if (task == null) return;
      this.head = taskConverter.fromBytes(task);
    }
    queueFile.remove();
    if (this.head != null) {
      int weight = this.head.weight();
      currentWeight.getAndUpdate(x -> x > weight ? x - weight : 0);
      this.head = null;
    }
  }

  @Override
  public int size() {
    return queueFile.size();
  }

  @Override
  public void close() throws IOException {
    queueFile.close();
  }

  @Nullable
  @Override
  public Long weight() {
    return currentWeight.get();
  }

  @Nullable
  @Override
  public Long getAvailableBytes() {
    return queueFile.storageBytes() - queueFile.usedBytes();
  }

  @Nonnull
  @Override
  public Iterator iterator() {
    Iterator iterator = queueFile.iterator();
    return new Iterator() {

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

      @Override
      public T next() {
        byte[] data = iterator.next();
        try {
          return taskConverter.fromBytes(data);
        } catch (IOException e) {
          throw Utils.throwAny(e);
        }
      }

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