com.axibase.tsd.client.PlainSender.java Source code

Java tutorial

Introduction

Here is the source code for com.axibase.tsd.client.PlainSender.java

Source

/*
 * Copyright 2015 Axibase Corporation or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 * https://www.axibase.com/atsd/axibase-apache-2.0.pdf
 *
 * or in the "license" file accompanying this file. This file 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.axibase.tsd.client;

import com.axibase.tsd.model.system.ClientConfiguration;
import com.axibase.tsd.plain.MarkerCommand;
import com.axibase.tsd.plain.PlainCommand;
import org.apache.commons.lang3.StringUtils;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.config.ConnectionConfig;
import org.apache.http.entity.AbstractHttpEntity;
import org.apache.http.entity.BufferedHttpEntity;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.impl.conn.PoolingHttpClientConnectionManager;
import org.glassfish.jersey.SslConfigurator;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.xml.bind.DatatypeConverter;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.*;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;

import static com.axibase.tsd.util.AtsdUtil.MARKER_KEYWORD;

/**
 * @author Nikolay Malevanny.
 */
class PlainSender extends AbstractHttpEntity implements Runnable {
    private static final Logger log = LoggerFactory.getLogger(PlainSender.class);
    private static final String PING_COMMAND = "ping\n";
    private static final int SMALL = 64;

    private String url;
    private CountDownLatch latch = new CountDownLatch(1);
    private CloseableHttpClient httpClient;
    private BlockingQueue<String> messages;
    private Map<String, List<String>> markerToMessages = new LinkedHashMap<String, List<String>>();
    private volatile SenderState state = SenderState.NEW;
    private final long pingTimeoutMillis;
    private long lastMessageTime;
    private CloseableHttpResponse response;
    private final ClientConfiguration clientConfiguration;
    private PoolingHttpClientConnectionManager connectionManager;

    public PlainSender(ClientConfiguration clientConfiguration, PlainSender old) {
        this.clientConfiguration = clientConfiguration;
        this.url = clientConfiguration.getDataUrl();
        this.pingTimeoutMillis = clientConfiguration.getPingTimeoutMillis();
        if (old != null) {
            messages = old.messages;
            markerToMessages = old.markerToMessages;
            log.info("Reborn plain commands sender using previous messages, size: {}", messages.size());
        }
    }

    public void send(PlainCommand plainCommand) {
        try {
            latch.await();
        } catch (InterruptedException e) {
            log.error("Initialization error:", e);
        }

        if (state != SenderState.WORKING) {
            throw new IllegalStateException("Could not send command using incorrect sender");
        }

        String text = plainCommand.compose();
        if (StringUtils.isBlank(text)) {
            log.error("Command is empty");
            return;
        }

        if (!text.endsWith("\n")) {
            text = text + "\n";
        }
        messages.add(text);
        log.debug("Message is added to queue, queue size = {}", (messages == null) ? 0 : messages.size());
    }

    @Override
    public boolean isRepeatable() {
        return false;
    }

    @Override
    public boolean isChunked() {
        return false;
    }

    @Override
    public long getContentLength() {
        return -1;
    }

    @Override
    public InputStream getContent() throws IOException, IllegalStateException {
        return null;
    }

    @Override
    public void writeTo(OutputStream outputStream) throws IOException {
        String marker = null;
        while (state == SenderState.WORKING) {
            String message = null;
            try {
                message = messages.poll(pingTimeoutMillis, TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                if (state != SenderState.WORKING) {
                    log.error("Could not poll message from queue", e);
                }
            }

            try {
                if (message != null) {
                    if (!clientConfiguration.isSkipStreamingControl()) {
                        if (marker == null && !message.startsWith(MARKER_KEYWORD)) {
                            MarkerCommand markerCommand = new MarkerCommand();
                            marker = markerCommand.getMarker();
                            write(outputStream, markerCommand.compose());
                        }
                    }

                    log.debug("Write message: {}", message);
                    write(outputStream, message);

                    if (!clientConfiguration.isSkipStreamingControl()) {
                        if (message.startsWith(MARKER_KEYWORD)) {
                            marker = StringUtils.removeStart(message, MARKER_KEYWORD).trim();
                            if (StringUtils.isBlank(marker)) {
                                throw new IllegalArgumentException("Bad marker message: " + message);
                            }
                        } else {
                            add(marker, message);
                        }
                    }

                    lastMessageTime = System.currentTimeMillis();
                }
            } catch (Throwable e) {
                log.error("Sender is broken, close it. Could not send message: {}", message, e);
                messages.add(message);
                close();
                return;
            }
            if (lastMessageTime - System.currentTimeMillis() > pingTimeoutMillis) {
                write(outputStream, PING_COMMAND);
                if (!clientConfiguration.isSkipStreamingControl()) {
                    add(marker, PING_COMMAND);
                }
                lastMessageTime = System.currentTimeMillis();
            }
        }
    }

    private void write(OutputStream outputStream, String text) throws IOException {
        outputStream.write(text.getBytes());
        outputStream.flush();
    }

    private void add(String marker, String message) {
        if (clientConfiguration.isSkipStreamingControl()) {
            log.error(
                    "Could not add message to marker, because streaming control is skipped, marker = {}, message= {}",
                    marker, message);
            throw new IllegalStateException("Could not add message to marker during skipped streaming control");
        }

        List<String> stored = markerToMessages.get(marker);
        if (stored == null) {
            stored = new ArrayList<String>();
            markerToMessages.put(marker, stored);
        }
        stored.add(message);
    }

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

    @Override
    public void run() {
        if (messages == null) {
            messages = new LinkedBlockingQueue<String>();
        }
        HttpPost httpPost = null;
        try {
            SslConfigurator sslConfig = SslConfigurator.newInstance().securityProtocol("SSL");
            connectionManager = HttpClient.createConnectionManager(clientConfiguration, sslConfig);
            connectionManager.setDefaultConnectionConfig(ConnectionConfig.custom().setBufferSize(SMALL).build());
            httpClient = HttpClients.custom().setConnectionManager(connectionManager).build();
            httpPost = new HttpPost(fullUrl());
            httpPost.setHeader("Authorization", "Basic " + DatatypeConverter.printBase64Binary(
                    (clientConfiguration.getUsername() + ":" + clientConfiguration.getPassword()).getBytes()));
            httpPost.setEntity(new BufferedHttpEntity(this));
        } catch (Throwable e) {
            log.error("Could not create http client: ", e);
            latch.countDown();
            close();
            return;
        }
        try {
            log.info("Start writing commands to {}", fullUrl());
            state = SenderState.WORKING;
            latch.countDown();
            response = httpClient.execute(httpPost);
        } catch (Throwable e) {
            log.error("Could not execute HTTP POST: {}", httpPost, e);
        } finally {
            log.warn("Http post execution is finished, close sender");
            close();
        }
    }

    public void close() {
        if (state == SenderState.CLOSED) {
            return;
        }
        log.info("Stop writing commands to {}", fullUrl());
        if (response != null) {
            try {
                response.close();
            } catch (IOException e) {
                log.error("Could not close response: {}", response, e);
            }
        }
        if (httpClient != null) {
            try {
                httpClient.close();
            } catch (IOException e) {
                log.error("Could not close client: {}", httpClient, e);
            }
        }
        if (connectionManager != null) {
            connectionManager.close();
        }
        state = SenderState.CLOSED;
    }

    private String fullUrl() {
        return url + "/command";
    }

    public boolean isWorking() {
        return state == SenderState.WORKING;
    }

    Map<String, List<String>> getMarkerToMessages() {
        return markerToMessages;
    }

    public boolean isClosed() {
        return state == SenderState.CLOSED;
    }

    private static enum SenderState {
        NEW, WORKING, CLOSED
    }
}