org.graylog2.telemetry.TelemetryPeriodical Maven / Gradle / Ivy

Go to download
package org.graylog2.telemetry;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.google.common.collect.EvictingQueue;
import com.google.common.collect.Lists;
import com.google.common.net.HttpHeaders;
import com.google.common.primitives.Ints;
import com.squareup.okhttp.Credentials;
import com.squareup.okhttp.Headers;
import com.squareup.okhttp.MediaType;
import com.squareup.okhttp.OkHttpClient;
import com.squareup.okhttp.Request;
import com.squareup.okhttp.RequestBody;
import com.squareup.okhttp.Response;
import org.graylog2.plugin.Version;
import org.graylog2.plugin.periodical.Periodical;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
import org.slf4j.Logger;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.net.URL;
import java.util.List;

import static com.google.common.base.Preconditions.checkNotNull;

public abstract class TelemetryPeriodical extends Periodical {
    private static final MediaType CONTENT_TYPE = MediaType.parse("application/x-jackson-smile");
    private static final String TELEMETRY_VERSION = TelemetryMetaData.VERSION.toString();
    private static final String USER_AGENT = "Graylog " + Version.CURRENT_CLASSPATH;

    protected final TelemetryConfiguration config;
    protected final EvictingQueue cachedRequestsQueue;
    protected final ObjectMapper objectMapper;
    protected final OkHttpClient httpClient;
    protected final String credentials;
    protected final URL telemetryUrl;
    protected final String filenamePattern;

    private final Logger log = getLogger();

    protected TelemetryPeriodical(TelemetryConfiguration config,
                                  EvictingQueue telemetryRequestsQueue,
                                  OkHttpClient httpClient,
                                  ObjectMapper objectMapper,
                                  URL telemetryUrl,
                                  String filenamePattern) {
        this.config = checkNotNull(config);
        this.cachedRequestsQueue = checkNotNull(telemetryRequestsQueue);
        this.httpClient = checkNotNull(httpClient);
        this.objectMapper = checkNotNull(objectMapper);
        this.credentials = Credentials.basic("telemetry", config.getToken());
        this.telemetryUrl = checkNotNull(telemetryUrl);
        this.filenamePattern = checkNotNull(filenamePattern);
    }

    protected abstract byte[] buildPayload();

    @Override
    public void doRun() {
        log.debug("Telemetry is activated: Transmitting node data sets.");
        final byte[] requestBody = buildPayload();

        if (config.isOfflineMode()) {
            final String filename = String.format(filenamePattern, DateTime.now(DateTimeZone.UTC).getMillis());

            if (!config.getDirectory().exists()) {
                boolean success = config.getDirectory().mkdirs();
                if (!success) {
                    log.error("Couldn't create directory {}", config.getDirectory().getAbsolutePath());
                    return;
                }
            }

            final File file = new File(config.getDirectory(), filename);

            log.debug("Telemetry is in offline mode. Writing data into {}", file);
            try (final OutputStream outputStream = new FileOutputStream(file)) {
                outputStream.write(requestBody);
            } catch (IOException e) {
                log.warn("Couldn't write telemetry data into " + file, e);
            }
        } else {
            final Headers headers = new Headers.Builder()
                    .add(HttpHeaders.AUTHORIZATION, credentials)
                    .add(HttpHeaders.USER_AGENT, USER_AGENT)
                    .add("X-Telemetry-Version", TELEMETRY_VERSION)
                    .build();

            final TelemetryRequest request = TelemetryRequest.create(headers, requestBody);

            boolean success = uploadDataSet(request);
            if (!success) {
                cachedRequestsQueue.add(request);
            } else if (!cachedRequestsQueue.isEmpty()) {
                log.debug("Trying to upload {} queued data sets", cachedRequestsQueue.size());
                uploadQueuedDataSets();
            }
        }
    }

    protected void uploadQueuedDataSets() {
        final List unsuccessfulRequests = Lists.newArrayListWithCapacity(cachedRequestsQueue.size());

        while (!cachedRequestsQueue.isEmpty()) {
            final TelemetryRequest request = cachedRequestsQueue.poll();
            boolean success = uploadDataSet(request);

            if (!success) {
                log.debug("Couldn't successfully upload telemetry data set, re-queueing data");
                unsuccessfulRequests.add(request);
            }
        }

        cachedRequestsQueue.addAll(unsuccessfulRequests);
    }

    protected boolean uploadDataSet(TelemetryRequest telemetryRequest) {
        final Request request = new Request.Builder()
                .url(telemetryUrl)
                .headers(telemetryRequest.headers())
                .post(RequestBody.create(CONTENT_TYPE, telemetryRequest.body()))
                .build();

        final Response response;
        try {
            response = httpClient.newCall(request).execute();
        } catch (IOException e) {
            log.error("Error while uploading telemetry data");
            log.debug("Error details", e);
            return false;
        }

        if (!response.isSuccessful()) {
            log.warn("Couldn't successfully upload telemetry data set: {}", response);
            return false;
        }

        return true;
    }

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

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

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

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

    @Override
    public int getInitialDelaySeconds() {
        return 60;
    }

    @Override
    public int getPeriodSeconds() {
        return Ints.saturatedCast(config.getReportInterval().toSeconds());
    }
}