/**
 * 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.alibaba.jstorm.daemon.supervisor;

import java.io.File;
import java.io.IOException;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.alibaba.jstorm.client.ConfigExtension;
import com.alibaba.jstorm.container.CgroupCenter;
import com.alibaba.jstorm.container.Hierarchy;
import com.alibaba.jstorm.container.SubSystemType;
import com.alibaba.jstorm.container.cgroup.CgroupCommon;
import com.alibaba.jstorm.container.cgroup.core.CgroupCore;
import com.alibaba.jstorm.container.cgroup.core.CpuCore;
import com.alibaba.jstorm.utils.JStormUtils;

/**
 * @author Johnfang (xiaojian.fxj@alibaba-inc.com)
 */
public class CgroupManager {
    private static final Logger LOG = LoggerFactory.getLogger(CgroupManager.class);

    public static final String JSTORM_HIERARCHY_NAME = "jstorm_cpu";
    public static final int ONE_CPU_SLOT = 1024;

    private CgroupCenter center;
    private Hierarchy h;
    private CgroupCommon rootCgroup;

    private static String rootDir;
    private static String cgroupBaseDir;

    public CgroupManager(Map conf) {
        LOG.info("running on cgroup mode");

        // Cgconfig service is used to create the corresponding cpu hierarchy "/cgroup/cpu"
        cgroupBaseDir = ConfigExtension.getCgroupBaseDir(conf);
        rootDir = ConfigExtension.getCgroupRootDir(conf);
        if (StringUtils.isBlank(rootDir) || StringUtils.isBlank(cgroupBaseDir)) {
            StringBuilder sb = new StringBuilder();
            sb.append("Check configuration file. The setting of ");
            sb.append(ConfigExtension.CGROUP_BASE_DIR).append(" and ");
            sb.append(ConfigExtension.CGROUP_ROOT_DIR);
            sb.append(" are invalid.");
            throw new RuntimeException(sb.toString());
        }

        String realRootDir = cgroupBaseDir + "/" + rootDir;
        File file = new File(realRootDir);
        if (!file.exists()) {
            LOG.error(realRootDir + " does not exist.");
            throw new RuntimeException("Check if cgconfig service starts or /etc/cgconfig.conf is " +
                    "consistent with configuration file.");
        }
        center = CgroupCenter.getInstance();
        if (center == null)
            throw new RuntimeException("Cgroup error, please check /proc/cgroups");
        this.prepareSubSystem();
    }

    private int validateCpuUpperLimitValue(int value) {
        /*
         * Valid value is -1 or 1~10 -1 means no control
         */
        if (value > 10)
            value = 10;
        else if (value < 1 && value != -1)
            value = 1;

        return value;
    }

    private void setCpuUsageUpperLimit(CpuCore cpuCore, int cpuCoreUpperLimit) throws IOException {
        /*
         * User cfs_period & cfs_quota to control the upper limit use of cpu core e.g. 
         * If making a process to fully use two cpu cores, set cfs_period_us to
         * 100000 and set cfs_quota_us to 200000 The highest value of "cpu core upper limit" is 10
         */
        cpuCoreUpperLimit = validateCpuUpperLimitValue(cpuCoreUpperLimit);

        if (cpuCoreUpperLimit == -1) {
            // No control of cpu usage
            cpuCore.setCpuCfsQuotaUs(cpuCoreUpperLimit);
        } else {
            cpuCore.setCpuCfsPeriodUs(100000);
            cpuCore.setCpuCfsQuotaUs(cpuCoreUpperLimit * 100000);
        }
    }

    public String startNewWorker(Map conf, int cpuNum, String workerId) throws SecurityException, IOException {
        CgroupCommon workerGroup = new CgroupCommon(workerId, h, this.rootCgroup);
        this.center.create(workerGroup);
        CgroupCore cpu = workerGroup.getCores().get(SubSystemType.cpu);
        CpuCore cpuCore = (CpuCore) cpu;
        cpuCore.setCpuShares(cpuNum * ONE_CPU_SLOT);
        setCpuUsageUpperLimit(cpuCore, ConfigExtension.getWorkerCpuCoreUpperLimit(conf));

        StringBuilder sb = new StringBuilder();
        sb.append("cgexec -g cpu:").append(workerGroup.getName()).append(" ");
        return sb.toString();
    }

    public void shutDownWorker(String workerId, boolean isKilled) {
        CgroupCommon workerGroup = new CgroupCommon(workerId, h, this.rootCgroup);
        try {
            if (!isKilled) {
                for (Integer pid : workerGroup.getTasks()) {
                    JStormUtils.kill(pid);
                }
                JStormUtils.sleepMs(1500);
            }
            center.delete(workerGroup);
        } catch (Exception e) {
            LOG.info("No task of " + workerId);
        }
    }

    public void close() throws IOException {
        this.center.delete(this.rootCgroup);
    }

    private void prepareSubSystem() {
        h = center.busy(SubSystemType.cpu);
        if (h == null) {
            Set<SubSystemType> types = new HashSet<>();
            types.add(SubSystemType.cpu);
            h = new Hierarchy(JSTORM_HIERARCHY_NAME, types, cgroupBaseDir);
        }
        rootCgroup = new CgroupCommon(rootDir, h, h.getRootCgroups());
    }
}