gobblin.writer.http.AbstractHttpWriter.java Source code

Java tutorial

Introduction

Here is the source code for gobblin.writer.http.AbstractHttpWriter.java

Source

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 gobblin.writer.http;

import java.io.IOException;
import java.net.URI;
import java.net.URISyntaxException;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;

import org.apache.commons.lang3.builder.ToStringBuilder;
import org.apache.commons.lang3.builder.ToStringStyle;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.conn.ConnectionRequest;
import org.apache.http.conn.HttpClientConnectionManager;
import org.apache.http.conn.routing.HttpRoute;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.util.EntityUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.base.Optional;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;

import gobblin.instrumented.writer.InstrumentedDataWriter;
import gobblin.util.ExecutorsUtils;

/**
 * Base class for HTTP writers. Defines the main extension points for different implementations.
 */
public abstract class AbstractHttpWriter<D> extends InstrumentedDataWriter<D> implements HttpWriterDecoration<D> {

    // Immutable state
    protected final Logger log;
    protected final boolean debugLogEnabled;
    protected final CloseableHttpClient client;
    private final ListeningExecutorService singleThreadPool;

    // Mutable state
    private URI curHttpHost = null;
    private long numRecordsWritten = 0L;
    private long numBytesWritten = 0L; //AbstractHttpWriter won't update as it could be expensive.
    Optional<HttpUriRequest> curRequest = Optional.absent();

    class HttpClientConnectionManagerWithConnTracking extends DelegatingHttpClientConnectionManager {

        public HttpClientConnectionManagerWithConnTracking(HttpClientConnectionManager fallback) {
            super(fallback);
        }

        @Override
        public ConnectionRequest requestConnection(HttpRoute route, Object state) {
            try {
                onConnect(new URI(route.getTargetHost().toURI()));
            } catch (IOException | URISyntaxException e) {
                throw new RuntimeException("onConnect() callback failure: " + e, e);
            }
            return super.requestConnection(route, state);
        }

    }

    @SuppressWarnings("rawtypes")
    public AbstractHttpWriter(AbstractHttpWriterBuilder builder) {
        super(builder.getState());
        this.log = builder.getLogger().isPresent() ? (Logger) builder.getLogger()
                : LoggerFactory.getLogger(this.getClass());
        this.debugLogEnabled = this.log.isDebugEnabled();

        HttpClientBuilder httpClientBuilder = builder.getHttpClientBuilder();
        httpClientBuilder.setConnectionManager(
                new HttpClientConnectionManagerWithConnTracking(builder.getHttpConnManager()));
        this.client = httpClientBuilder.build();
        this.singleThreadPool = MoreExecutors.listeningDecorator(Executors.newSingleThreadExecutor());

        if (builder.getSvcEndpoint().isPresent()) {
            setCurServerHost((URI) builder.getSvcEndpoint().get());
        }
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void cleanup() throws IOException {
        this.client.close();
        ExecutorsUtils.shutdownExecutorService(this.singleThreadPool, Optional.of(log));
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public void close() throws IOException {
        cleanup();
        super.close();
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public long recordsWritten() {
        return this.numRecordsWritten;
    }

    /**
     * {@inheritDoc}
     */
    @Override
    public long bytesWritten() throws IOException {
        return this.numBytesWritten;
    }

    /**
     * Send and process the request. If it's a retry request, skip onNewRecord method call and go straight sending request.
     * {@inheritDoc}
     */
    @Override
    public void writeImpl(D record) throws IOException {
        if (!isRetry()) {
            //If currentRequest is still here, it means this is retry request.
            //In this case, don't invoke onNewRecord again as onNewRecord is not guaranteed to be idempotent.
            //(e.g: If you do batch processing duplicate record can go in, etc.)
            curRequest = onNewRecord(record);
        }

        if (curRequest.isPresent()) {
            ListenableFuture<CloseableHttpResponse> responseFuture = sendRequest(curRequest.get());
            try (CloseableHttpResponse response = waitForResponse(responseFuture)) {
                processResponse(response);
            }
            curRequest = Optional.absent(); //Clear request if successful
        }
        numRecordsWritten++;
    }

    /**
     * Prior to commit, it will invoke flush method to flush any remaining item if writer uses batch
     * {@inheritDoc}
     * @see gobblin.instrumented.writer.InstrumentedDataWriterBase#commit()
     */
    @Override
    public void commit() throws IOException {
        flush();
        super.commit();
    }

    /**
     * If writer supports batch, override this method.
     * (Be aware of failure and retry as flush can be called multiple times in case of failure @see SalesforceRestWriter )
     */
    public void flush() {
    }

    /**
     * Sends request using single thread pool so that it can be easily terminated(use case: time out)
     * {@inheritDoc}
     * @see gobblin.writer.http.HttpWriterDecoration#sendRequest(org.apache.http.client.methods.HttpUriRequest)
     */
    @Override
    public ListenableFuture<CloseableHttpResponse> sendRequest(final HttpUriRequest request) throws IOException {
        return singleThreadPool.submit(new Callable<CloseableHttpResponse>() {
            @Override
            public CloseableHttpResponse call() throws Exception {
                return client.execute(request);
            }
        });
    }

    /**
     * Checks if it's retry request.
     * All successful request should make currentRequest absent. If currentRequest still exists, it means there was a failure.
     * There's couple of methods need this indicator such as onNewRecord, since it is not a new record.
     * @return true if current request it holds is retry.
     */
    public boolean isRetry() {
        return curRequest.isPresent();
    }

    /**
     * Default implementation is to use HttpClients socket timeout which is waiting based on elapsed time between
     * last packet sent from client till receive it from server.
     *
     * {@inheritDoc}
     * @see gobblin.writer.http.HttpWriterDecoration#waitForResponse(com.google.common.util.concurrent.ListenableFuture)
     */
    @Override
    public CloseableHttpResponse waitForResponse(ListenableFuture<CloseableHttpResponse> responseFuture) {
        try {
            return responseFuture.get();
        } catch (InterruptedException | ExecutionException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * Default implementation where any status code equal to or greater than 400 is regarded as a failure.
     * {@inheritDoc}
     * @see gobblin.writer.http.HttpWriterDecoration#processResponse(org.apache.http.HttpResponse)
     */
    @Override
    public void processResponse(CloseableHttpResponse response) throws IOException, UnexpectedResponseException {
        if (response.getStatusLine().getStatusCode() >= 400) {
            if (response.getEntity() != null) {
                throw new RuntimeException("Failed. " + EntityUtils.toString(response.getEntity()) + " , response: "
                        + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE));
            }
            throw new RuntimeException("Failed. Response: "
                    + ToStringBuilder.reflectionToString(response, ToStringStyle.SHORT_PREFIX_STYLE));
        }
    }

    public Logger getLog() {
        return this.log;
    }

    public URI getCurServerHost() {
        if (null == this.curHttpHost) {
            setCurServerHost(chooseServerHost());
        }
        if (null == this.curHttpHost) {
            throw new RuntimeException("No server host selected!");
        }
        return this.curHttpHost;
    }

    /** Clears the current http host so that next request will trigger a new selection using
     * {@link #chooseServerHost() */
    void clearCurServerHost() {
        this.curHttpHost = null;
    }

    void setCurServerHost(URI curHttpHost) {
        this.log.info("Setting current HTTP server host to: " + curHttpHost);
        this.curHttpHost = curHttpHost;
    }

}