org.springframework.xd.greenplum.gpfdist.GPFDistMessageHandler.java Source code

Introduction

Here is the source code for org.springframework.xd.greenplum.gpfdist.GPFDistMessageHandler.java

Source

/*
 * Copyright 2015 the original author or authors.
 *
 * 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 org.springframework.xd.greenplum.gpfdist;

import java.util.Date;
import java.util.concurrent.FutureTask;
import java.util.concurrent.TimeUnit;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.reactivestreams.Processor;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHandlingException;
import org.springframework.scheduling.TaskScheduler;
import org.springframework.util.StringUtils;
import org.springframework.util.concurrent.SettableListenableFuture;
import org.springframework.xd.greenplum.support.GreenplumLoad;
import org.springframework.xd.greenplum.support.NetworkUtils;
import org.springframework.xd.greenplum.support.RuntimeContext;

import reactor.Environment;
import reactor.core.processor.RingBufferProcessor;
import reactor.io.buffer.Buffer;

import com.codahale.metrics.Meter;

public class GPFDistMessageHandler extends AbstractGPFDistMessageHandler {

    private final Log log = LogFactory.getLog(GPFDistMessageHandler.class);

    private final int port;

    private final int flushCount;

    private final int flushTime;

    private final int batchTimeout;

    private final int batchCount;

    private final int batchPeriod;

    private final String delimiter;

    private GreenplumLoad greenplumLoad;

    private Processor<Buffer, Buffer> processor;

    private GPFDistServer gpfdistServer;

    private TaskScheduler sqlTaskScheduler;

    private final TaskFuture taskFuture = new TaskFuture();

    private int rateInterval = 0;
    private Meter meter = null;
    private int meterCount = 0;

    public GPFDistMessageHandler(int port, int flushCount, int flushTime, int batchTimeout, int batchCount,
            int batchPeriod, String delimiter) {
        super();
        this.port = port;
        this.flushCount = flushCount;
        this.flushTime = flushTime;
        this.batchTimeout = batchTimeout;
        this.batchCount = batchCount;
        this.batchPeriod = batchPeriod;
        this.delimiter = StringUtils.hasLength(delimiter) ? delimiter : null;
    }

    @Override
    protected void doWrite(Message<?> message) throws Exception {
        Object payload = message.getPayload();
        if (payload instanceof String) {
            String data = (String) payload;
            if (delimiter != null) {
                processor.onNext(Buffer.wrap(data + delimiter));
            } else {
                processor.onNext(Buffer.wrap(data));
            }
            if (meter != null) {
                if ((meterCount++ % rateInterval) == 0) {
                    meter.mark(rateInterval);
                    log.info("METER: 1 minute rate = " + meter.getOneMinuteRate() + " mean rate = "
                            + meter.getMeanRate());
                }
            }
        } else {
            throw new MessageHandlingException(message, "message not a String");
        }
    }

    @Override
    protected void onInit() throws Exception {
        super.onInit();
        Environment.initializeIfEmpty().assignErrorJournal();
        processor = RingBufferProcessor.create(false);
    }

    @Override
    protected void doStart() {
        try {
            log.info("Creating gpfdist protocol listener on port=" + port);
            gpfdistServer = new GPFDistServer(processor, port, flushCount, flushTime, batchTimeout, batchCount);
            gpfdistServer.start();
            log.info("gpfdist protocol listener running on port=" + gpfdistServer.getLocalPort());
        } catch (Exception e) {
            throw new RuntimeException("Error starting protocol listener", e);
        }

        if (greenplumLoad != null) {
            log.info("Scheduling gpload task with batchPeriod=" + batchPeriod);

            final RuntimeContext context = new RuntimeContext();
            context.addLocation(NetworkUtils.getGPFDistUri(gpfdistServer.getLocalPort()));

            sqlTaskScheduler.schedule((new FutureTask<Void>(new Runnable() {
                @Override
                public void run() {
                    boolean taskValue = true;
                    try {
                        while (!taskFuture.interrupted) {
                            try {
                                greenplumLoad.load(context);
                            } catch (Exception e) {
                                log.error("Error in load", e);
                            }
                            Thread.sleep(batchPeriod * 1000);
                        }
                    } catch (Exception e) {
                        taskValue = false;
                    }
                    taskFuture.set(taskValue);
                }
            }, null)), new Date());

        } else {
            log.info("Skipping gpload tasks because greenplumLoad is not set");
        }
    }

    @Override
    protected void doStop() {
        if (greenplumLoad != null) {
            taskFuture.interruptTask();
            try {
                long now = System.currentTimeMillis();
                // wait a bit more than batch period
                Boolean value = taskFuture.get(batchTimeout + batchPeriod + 2, TimeUnit.SECONDS);
                log.info("Stopping, got future value " + value + " from task which took "
                        + (System.currentTimeMillis() - now) + "ms");
            } catch (Exception e) {
                log.warn("Got error from task wait value which may indicate trouble", e);
            }
        }

        try {
            processor.onComplete();
            gpfdistServer.stop();
        } catch (Exception e) {
            log.warn("Error shutting down protocol listener", e);
        }
    }

    public void setSqlTaskScheduler(TaskScheduler sqlTaskScheduler) {
        this.sqlTaskScheduler = sqlTaskScheduler;
    }

    public void setGreenplumLoad(GreenplumLoad greenplumLoad) {
        this.greenplumLoad = greenplumLoad;
    }

    public void setRateInterval(int rateInterval) {
        this.rateInterval = rateInterval;
        if (rateInterval > 0) {
            meter = new Meter();
        }
    }

    private static class TaskFuture extends SettableListenableFuture<Boolean> {

        boolean interrupted = false;

        @Override
        protected void interruptTask() {
            interrupted = true;
        }
    }
}