com.linkedin.pinot.controller.validation.StorageQuotaChecker.java Source code

Introduction

Here is the source code for com.linkedin.pinot.controller.validation.StorageQuotaChecker.java

Source

/**
 * Copyright (C) 2014-2016 LinkedIn Corp. (pinot-core@linkedin.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.linkedin.pinot.controller.validation;

import java.io.File;
import org.apache.commons.io.FileUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.google.common.base.Preconditions;
import com.linkedin.pinot.common.config.AbstractTableConfig;
import com.linkedin.pinot.common.config.QuotaConfig;
import com.linkedin.pinot.controller.api.restlet.resources.TableSizeReader;
import javax.annotation.Nonnegative;
import javax.annotation.Nonnull;

/**
 * Class to check if a new segment is within the configured storage quota for the table
 *
 */
public class StorageQuotaChecker {
    private static final Logger LOGGER = LoggerFactory.getLogger(StorageQuotaChecker.class);

    private final TableSizeReader tableSizeReader;
    private final AbstractTableConfig tableConfig;

    public StorageQuotaChecker(AbstractTableConfig tableConfig, TableSizeReader tableSizeReader) {
        this.tableConfig = tableConfig;
        this.tableSizeReader = tableSizeReader;
    }

    public class QuotaCheckerResponse {
        public boolean isSegmentWithinQuota;
        public String reason;

        QuotaCheckerResponse(boolean isSegmentWithinQuota, String reason) {
            this.isSegmentWithinQuota = isSegmentWithinQuota;
            this.reason = reason;
        }
    }

    /**
     * check if the segment represented by segmentFile is within the storage quota
     * @param segmentFile untarred segment. This should not be null.
     *                    segmentFile must exist on disk and must be a directory
     * @param tableNameWithType table name without type (OFFLINE/REALTIME) information
     * @param segmentName name of the segment being added
     * @param timeoutMsec timeout in milliseconds for reading table sizes from server
     *
     */
    public QuotaCheckerResponse isSegmentStorageWithinQuota(@Nonnull File segmentFile,
            @Nonnull String tableNameWithType, @Nonnull String segmentName, @Nonnegative int timeoutMsec) {
        Preconditions.checkNotNull(segmentFile);
        Preconditions.checkNotNull(tableNameWithType);
        Preconditions.checkNotNull(segmentName);
        Preconditions.checkArgument(timeoutMsec > 0, "Timeout value must be > 0, input: %s", timeoutMsec);
        Preconditions.checkArgument(segmentFile.exists(), "Segment file: %s does not exist", segmentFile);
        Preconditions.checkArgument(segmentFile.isDirectory(), "Segment file: %s is not a directory", segmentFile);

        // 1. Read table config
        // 2. read table size from all the servers
        // 3. update predicted segment sizes
        // 4. is the updated size within quota
        QuotaConfig quotaConfig = tableConfig.getQuotaConfig();
        int numReplicas = tableConfig.getValidationConfig().getReplicationNumber();
        final String tableName = tableConfig.getTableName();

        if (quotaConfig == null) {
            // no quota configuration...so ignore for backwards compatibility
            return new QuotaCheckerResponse(true, "Quota configuration not set for table: " + tableNameWithType);
        }

        long allowedStorageBytes = numReplicas * quotaConfig.storageSizeBytes();
        if (allowedStorageBytes < 0) {
            return new QuotaCheckerResponse(true,
                    "Storage quota is not configured for table: " + tableNameWithType);
        }

        long incomingSegmentSizeBytes = FileUtils.sizeOfDirectory(segmentFile);

        // read table size
        TableSizeReader.TableSubTypeSizeDetails tableSubtypeSize = tableSizeReader
                .getTableSubtypeSize(tableNameWithType, timeoutMsec);

        // If the segment exists(refresh), get the existing size
        TableSizeReader.SegmentSizeDetails sizeDetails = tableSubtypeSize.segments.get(segmentName);
        long existingSegmentSizeBytes = sizeDetails != null ? sizeDetails.estimatedSizeInBytes : 0;

        long estimatedFinalSizeBytes = tableSubtypeSize.estimatedSizeInBytes - existingSegmentSizeBytes
                + incomingSegmentSizeBytes;
        if (estimatedFinalSizeBytes <= allowedStorageBytes) {
            return new QuotaCheckerResponse(true, String.format(
                    "Estimated size: %d bytes is within the configured quota of %d (bytes) for table %s. Incoming segment size: %d (bytes)",
                    estimatedFinalSizeBytes, allowedStorageBytes, tableName, incomingSegmentSizeBytes));
        } else {
            return new QuotaCheckerResponse(false, String.format(
                    "Estimated size: %d bytes exceeds the configured quota of %d (bytes) for table %s. Incoming segment size: %d (bytes)",
                    estimatedFinalSizeBytes, allowedStorageBytes, tableName, incomingSegmentSizeBytes));
        }
    }
}