com.aerospike.kafka.connect.sink.AsyncWriter Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of aerospike-kafka-connector
Aerospike Connector for Kafka Connect
The newest version!
/*
 * Copyright 2016-2017 Aerospike, Inc.
 *
 * Portions may be licensed to Aerospike, Inc. under one or more contributor
 * license agreements WHICH ARE COMPATIBLE WITH THE APACHE LICENSE, VERSION 2.0.
 *
 * 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 com.aerospike.kafka.connect.sink;

import org.apache.kafka.connect.errors.ConnectException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.aerospike.client.AerospikeClient;
import com.aerospike.client.AerospikeException;
import com.aerospike.client.Bin;
import com.aerospike.client.Host;
import com.aerospike.client.Key;
import com.aerospike.client.async.EventLoops;
import com.aerospike.client.async.NettyEventLoops;
import com.aerospike.client.policy.ClientPolicy;
import com.aerospike.client.policy.RecordExistsAction;
import com.aerospike.client.policy.WritePolicy;
import com.aerospike.kafka.connect.data.AerospikeRecord;

import io.netty.channel.EventLoopGroup;
import io.netty.channel.nio.NioEventLoopGroup;

/**
 * The AsyncWriter handles connections to the Aerospike cluster, sending data
 * and flush. The write sends individual request to write each record using the
 * asynchronous client. The flush method waits until all in-flight request have been
 * completed.
 */
public class AsyncWriter {

    private static final Logger log = LoggerFactory.getLogger(AsyncWriter.class);

    private final EventLoops eventLoops;
    private final AerospikeClient client;
    private final WritePolicy writePolicy;
    private final Counter inFlight;
    private final ResultListener listener;

    public AsyncWriter(ConnectorConfig config) {
        eventLoops = createEventLoops();
        client = connectClient(config, eventLoops);
        writePolicy = createWritePolicy(config);
        inFlight = new Counter();
        listener = new ResultListener(inFlight);
    }

    public void write(AerospikeRecord record) {
        listener.raiseErrors();
        Key key = record.key();
        Bin[] bins = record.bins();
        inFlight.increment();
        log.info("Writing record asynchronously: {}", key);
        client.put(eventLoops.next(), listener, writePolicy, key, bins);
    }

    public void flush() {
        listener.raiseErrors();
        inFlight.waitUntilZero();
    }

    public void close() {
        client.close();
        eventLoops.close();
    }

    private EventLoops createEventLoops() {
        EventLoopGroup group = new NioEventLoopGroup();
        return new NettyEventLoops(group);
    }

    private AerospikeClient connectClient(ConnectorConfig config, EventLoops eventLoops) {
        try {
            Host[] hosts = config.getHosts();
            ClientPolicy policy = new ClientPolicy();
            policy.eventLoops = eventLoops;
            AerospikeClient client = new AerospikeClient(policy, hosts);
            log.info("Connected to Aerospike cluster at {}", (Object[]) hosts);
            return client;
        } catch (AerospikeException e) {
            eventLoops.close();
            throw new ConnectException("Error connecting to Aerospike cluster", e);
        }
    }

    private WritePolicy createWritePolicy(ConnectorConfig config) {
        WritePolicy policy = new WritePolicy();
        RecordExistsAction action = config.getPolicyRecordExistsAction();
        if (action != null) {
            policy.recordExistsAction = action;
        }
        log.trace("Write Policy: recordExistsAction={}", policy.recordExistsAction);
        return policy;
    }
}