com.github.brandtg.switchboard.LogPuller.java Source code

Introduction

Here is the source code for com.github.brandtg.switchboard.LogPuller.java

Source

/**
 * Copyright (C) 2015 Greg Brandt (brandt.greg@gmail.com)
 *
 * 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.github.brandtg.switchboard;

import com.fasterxml.jackson.databind.ObjectMapper;
import org.apache.http.HttpEntity;
import org.apache.http.HttpHost;
import org.apache.http.HttpResponse;
import org.apache.http.client.HttpClient;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.impl.client.HttpClients;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.InputStream;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.URI;
import java.net.URLEncoder;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

public class LogPuller implements Runnable {
    private static final Logger LOG = LoggerFactory.getLogger(LogPuller.class);
    private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
    private static final String ENCODING = "UTF-8";

    private final String collection;
    private final long lastIndex;
    private final InetSocketAddress sourceAddress;
    private final InetSocketAddress sinkAddress;
    private final AtomicBoolean isShutdown;

    /**
     * Pulls logs for a collection from a switchboard server.
     *
     * @param sourceAddress
     *  The address of the switchboard server.
     * @param sinkAddress
     *  The address to provide as "target" in HTTP requests, on which to listen for data locally.
     * @param collection
     *  The name of the log.
     * @param lastIndex
     *  The last index consumed (-1 means nothing consumed).
     */
    public LogPuller(InetSocketAddress sourceAddress, InetSocketAddress sinkAddress, String collection,
            long lastIndex) {
        this.sourceAddress = sourceAddress;
        this.sinkAddress = sinkAddress;
        this.collection = collection;
        this.lastIndex = lastIndex;
        this.isShutdown = new AtomicBoolean(false);
    }

    /** Shuts down this log puller. */
    public void shutdown() {
        if (!isShutdown.getAndSet(true)) {
            LOG.info("Shut down log puller");
        }
    }

    @Override
    public void run() {
        HttpClient httpClient = HttpClients.createDefault();
        AtomicLong currentIndex = new AtomicLong(lastIndex);
        HttpHost host = new HttpHost(sourceAddress.getAddress(), sourceAddress.getPort());
        boolean firstLoop = true;

        while (!isShutdown.get()) {
            // Build URI
            StringBuilder sb = new StringBuilder();
            try {
                if (firstLoop) {
                    sb.append("/log/metadata/header?target=").append(sinkAddress.getHostName()).append(":")
                            .append(sinkAddress.getPort());
                    firstLoop = false;
                } else {
                    sb.append("/log/").append(URLEncoder.encode(collection, ENCODING)).append("/")
                            .append(currentIndex.get()).append("?target=").append(sinkAddress.getHostName())
                            .append(":").append(sinkAddress.getPort());
                }
            } catch (Exception e) {
                throw new RuntimeException(e);
            }

            // TODO: Clean this up a little bit
            int resStatus = -1;
            synchronized (this) {
                HttpEntity entity = null;
                try {
                    // Get data
                    URI uri = URI.create(sb.toString());
                    HttpGet req = new HttpGet(uri);
                    HttpResponse res = httpClient.execute(host, req);
                    entity = res.getEntity();
                    resStatus = res.getStatusLine().getStatusCode();
                    if (resStatus == 200) {
                        // Wait for data to be consumed
                        // n.b. This object will be registered as a listener
                        wait();

                        // Update position
                        InputStream inputStream = res.getEntity().getContent();
                        LogRegionResponse metadata = OBJECT_MAPPER.readValue(inputStream, LogRegionResponse.class);
                        currentIndex
                                .set(metadata.getLogRegions().get(metadata.getLogRegions().size() - 1).getIndex());

                        for (LogRegion logRegion : metadata.getLogRegions()) {
                            LOG.info("Received {}", logRegion);
                        }
                    }
                } catch (Exception e) {
                    LOG.error("Error", e);
                } finally {
                    if (entity != null) {
                        try {
                            EntityUtils.consume(entity);
                        } catch (IOException e) {
                            LOG.error("Error", e);
                        }
                    }
                }
            }

            // Sleep if did not get data
            if (resStatus != 200) {
                try {
                    LOG.debug("No data available, sleeping 1000 ms");
                    Thread.sleep(1000);
                } catch (InterruptedException e) {
                    LOG.warn("Error while sleeping for more data", e);
                }
            }
        }
    }
}