io.radiowitness.kinesis.producer.PutRecordTask.java Source code

Introduction

Here is the source code for io.radiowitness.kinesis.producer.PutRecordTask.java

Source

/*
 * Copyright (C) 2016 An Honest Effort LLC.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

package io.radiowitness.kinesis.producer;

import com.amazonaws.handlers.AsyncHandler;
import com.amazonaws.services.kinesis.AmazonKinesisAsyncClient;
import com.amazonaws.services.kinesis.model.PutRecordRequest;
import com.amazonaws.services.kinesis.model.PutRecordResult;
import com.google.common.util.concurrent.SettableFuture;
import io.radiowitness.proto.pack.MessagePacker;
import io.radiowitness.proto.pack.MessagePackingException;

import java.nio.ByteBuffer;
import java.util.Optional;
import java.util.TimerTask;

import static io.radiowitness.proto.Proto.BaseMessage;

public class PutRecordTask extends TimerTask implements AsyncHandler<PutRecordRequest, PutRecordResult> {

    public static final int PUT_PAYLOAD_SIZE = 25000;

    private final KinesisProducerConfig config;
    private final AmazonKinesisAsyncClient client;
    private final SettableFuture<String> future;
    private final String partitionKey;
    private final Optional<String> sequenceNumber;

    private final ByteBuffer packed;
    private final MessagePacker packer;
    private final Object txnLock = new Object();

    private boolean hasRun = false;

    public PutRecordTask(KinesisProducerConfig config, AmazonKinesisAsyncClient client,
            SettableFuture<String> future, String partitionKey, Optional<String> sequenceNumber) {
        this.config = config;
        this.client = client;
        this.future = future;
        this.partitionKey = partitionKey;
        this.sequenceNumber = sequenceNumber;
        packed = ByteBuffer.allocate(PUT_PAYLOAD_SIZE * config.getPayloadsPerRecordMax());
        packer = new MessagePacker(packed);
    }

    public SettableFuture<String> getFuture() {
        return future;
    }

    public boolean packOrRun(BaseMessage message) throws MessagePackingException {
        boolean packed;
        boolean wasFull;

        synchronized (txnLock) {
            if (hasRun) {
                return false;
            } else {
                wasFull = packer.isFull();
                packed = packer.pack(message);
            }
        }

        if (!wasFull && packer.isFull()) {
            super.cancel();
            this.run();
        }

        return packed;
    }

    @Override
    public void run() {
        synchronized (txnLock) {
            if (!hasRun) {
                packed.flip();

                if (packed.remaining() < 1) {
                    future.setException(new MessagePackingException("buffer is empty"));
                } else {
                    PutRecordRequest request = new PutRecordRequest();
                    request.setStreamName(config.getStreamName());
                    request.setPartitionKey(partitionKey);
                    request.setData(packed);

                    if (sequenceNumber.isPresent()) {
                        request.setSequenceNumberForOrdering(sequenceNumber.get());
                    }

                    client.putRecordAsync(request, this);
                    hasRun = true;
                }
            }
        }
    }

    @Override
    public void onSuccess(PutRecordRequest request, PutRecordResult result) {
        future.set(result.getSequenceNumber());
    }

    @Override
    public void onError(Exception e) {
        future.setException(e);
    }

}