com.pinterest.pinlater.PinLaterQueueConfig.java Source code

Java tutorial

Introduction

Here is the source code for com.pinterest.pinlater.PinLaterQueueConfig.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 com.pinterest.pinlater;

import com.pinterest.pinlater.commons.config.ConfigFileServerSet;
import com.pinterest.pinlater.commons.config.ConfigFileWatcher;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableSet;
import com.twitter.common.base.MorePreconditions;
import com.twitter.common.net.pool.DynamicHostSet;
import com.twitter.common.zookeeper.ServerSet;
import com.twitter.thrift.ServiceInstance;
import com.twitter.util.ExceptionalFunction;
import org.apache.commons.configuration.PropertiesConfiguration;
import org.codehaus.jackson.annotate.JsonIgnoreProperties;
import org.codehaus.jackson.map.ObjectMapper;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * Class that keeps track of per-queue dynamic configuration for PinLater.
 */
public class PinLaterQueueConfig {

    private static final Logger LOG = LoggerFactory.getLogger(PinLaterQueueConfig.class);

    private final AtomicReference<QueueConfigSchema> queueConfigSchemaRef = new AtomicReference<QueueConfigSchema>();
    private final AtomicInteger numPinLaterServers = new AtomicInteger(1);
    private final AtomicReference<ImmutableMap<String, QueueRateLimiter.IFace>> queueRateLimitMapRef = new AtomicReference<ImmutableMap<String, QueueRateLimiter.IFace>>();

    private final String queueConfigFilePath;
    private final String pinlaterServerSetPath;
    private final boolean pinlaterServerSetEnabled;

    public PinLaterQueueConfig(PropertiesConfiguration configuration) {
        this(configuration.getString("QUEUE_CONFIG_FILE_PATH"), configuration.getString("SERVER_SET_PATH"),
                configuration.getBoolean("SERVER_SET_ENABLED"));
    }

    @VisibleForTesting
    public PinLaterQueueConfig(String queueConfigFilePath, String pinlaterServerSetPath,
            boolean pinlaterServerSetEnabled) {
        this.queueConfigFilePath = queueConfigFilePath;
        this.pinlaterServerSetPath = pinlaterServerSetPath;
        this.pinlaterServerSetEnabled = pinlaterServerSetEnabled;
    }

    public void initialize() throws Exception {
        // Check if use of serverset is enabled, and if so, register a change monitor so we
        // can find out how many PinLater servers are active. We use this to compute the
        // per-server rate limit.
        if (pinlaterServerSetEnabled) {
            MorePreconditions.checkNotBlank(pinlaterServerSetPath);
            LOG.info("Monitoring pinlater serverset: {}", pinlaterServerSetPath);
            String fullServerSetPath = getClass().getResource("/" + pinlaterServerSetPath).getPath();
            ServerSet serverSet = new ConfigFileServerSet(fullServerSetPath);
            serverSet.monitor(new DynamicHostSet.HostChangeMonitor<ServiceInstance>() {
                @Override
                public void onChange(ImmutableSet<ServiceInstance> hostSet) {
                    int oldSize = numPinLaterServers.get();
                    int newSize = hostSet.size();
                    if (newSize == 0) {
                        LOG.error("PinLater serverset is empty, ignoring and keeping old size: {}", oldSize);
                        return;
                    }
                    if (oldSize == newSize) {
                        LOG.info("PinLater serverset update, size unchanged: {}", oldSize);
                        return;
                    }

                    LOG.info("PinLater serverset update, old size: {}, new size: {}", oldSize, newSize);
                    numPinLaterServers.set(newSize);
                    rebuild();
                }
            });
        } else {
            LOG.info("PinLater server set is disabled; rate limits will be applied per server.");
        }

        if (queueConfigFilePath == null || queueConfigFilePath.isEmpty()) {
            LOG.info("Queue config zookeeper path not specified, using defaults.");
            return;
        }

        LOG.info("Registering watch on queue config: {}", queueConfigFilePath);
        String fullQueueConfigFilePath = getClass().getResource("/" + queueConfigFilePath).getPath();
        ConfigFileWatcher.defaultInstance().addWatch(fullQueueConfigFilePath,
                new ExceptionalFunction<byte[], Void>() {
                    @Override
                    public Void applyE(byte[] bytes) throws Exception {
                        QueueConfigSchema queueConfigSchema = QueueConfigSchema.load(bytes);
                        LOG.info("Queue config update, new value: {}", queueConfigSchema);
                        queueConfigSchemaRef.set(queueConfigSchema);
                        rebuild();
                        return null;
                    }
                });
    }

    /**
     * Determines whether a dequeue request should be allowed.
     *
     * @param queueName  name of the queue.
     * @param numJobs    number of jobs intended to be dequeued.
     * @return whether to allow the request.
     */
    public boolean allowDequeue(String queueName, int numJobs) {
        MorePreconditions.checkNotBlank(queueName);
        Preconditions.checkArgument(numJobs > 0);
        ImmutableMap<String, QueueRateLimiter.IFace> queueRateLimitMap = queueRateLimitMapRef.get();
        if (queueRateLimitMap != null && queueRateLimitMap.containsKey(queueName)) {
            return queueRateLimitMap.get(queueName).allowDequeue(numJobs);
        } else {
            // No rate limit specified for this queue, so always allow.
            return true;
        }
    }

    @VisibleForTesting
    double getDequeueRate(String queueName) {
        MorePreconditions.checkNotBlank(queueName);
        ImmutableMap<String, QueueRateLimiter.IFace> queueRateLimitMap = queueRateLimitMapRef.get();
        if (queueRateLimitMap != null && queueRateLimitMap.containsKey(queueName)) {
            return queueRateLimitMap.get(queueName).getRate();
        } else {
            // No rate limit specified for this queue.
            return Double.MAX_VALUE;
        }
    }

    /**
     * Rebuilds the per-queue configuration. This method is thread-safe.
     */
    @VisibleForTesting
    synchronized void rebuild() {
        QueueConfigSchema queueConfigSchema = queueConfigSchemaRef.get();
        int numServers = numPinLaterServers.get();

        if (queueConfigSchema == null || queueConfigSchema.queues == null || queueConfigSchema.queues.isEmpty()) {
            return;
        }

        ImmutableMap.Builder<String, QueueRateLimiter.IFace> builder = new ImmutableMap.Builder<String, QueueRateLimiter.IFace>();
        for (QueueConfigSchema.Queue queue : queueConfigSchema.queues) {
            builder.put(queue.name, QueueRateLimiter.create(queue.queueConfig.maxJobsPerSecond / numServers));
        }

        queueRateLimitMapRef.set(builder.build());
    }

    @VisibleForTesting
    AtomicReference<QueueConfigSchema> getQueueConfigSchemaRef() {
        return queueConfigSchemaRef;
    }

    @VisibleForTesting
    AtomicInteger getNumPinLaterServers() {
        return numPinLaterServers;
    }

    /**
     * Defines the queue configuration json schema.
     */
    public static class QueueConfigSchema {

        private static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();

        public List<Queue> queues;

        public static class Queue {

            public String name;
            public QueueConfig queueConfig;
        }

        @JsonIgnoreProperties(ignoreUnknown = true)
        public static class QueueConfig {

            public double maxJobsPerSecond;
        }

        public static QueueConfigSchema load(byte[] bytes) throws IOException {
            return QueueConfigSchema.OBJECT_MAPPER.readValue(bytes, QueueConfigSchema.class);
        }

        @Override
        public String toString() {
            StringBuilder out = new StringBuilder();
            out.append("\nQueue Config:");
            for (Queue q : queues) {
                out.append("\nQueue: ").append(q.name);
                out.append(" maxJobsPerSecond: ").append(q.queueConfig.maxJobsPerSecond);
            }
            return out.toString();
        }
    }
}