text: // Copyright 2020 The Chromium Authors. All rights reserved.
  // Use of this source code is governed by a BSD-style license that can be
  // found in the LICENSE file.
  
  #include "base/cpu_affinity_posix.h"
  
  #include <sched.h>
  
  #include "base/cpu.h"
  #include "base/process/internal_linux.h"
  #include "third_party/abseil-cpp/absl/types/optional.h"
  
  namespace base {
  
  namespace {
  
  const cpu_set_t& AllCores() {
    static const cpu_set_t kAllCores = []() {
      cpu_set_t set;
      CPU_ZERO(&set);
      const std::vector<CPU::CoreType>& core_types = CPU::GetGuessedCoreTypes();
      if (core_types.empty()) {
        memset(&set, 0xff, sizeof(set));
      } else {
        for (size_t index = 0; index < core_types.size(); index++)
          CPU_SET(index, &set);
      }
      return set;
    }();
    return kAllCores;
  }
  
  const cpu_set_t& LittleCores() {
    static const cpu_set_t kLittleCores = []() {
      const std::vector<CPU::CoreType>& core_types = CPU::GetGuessedCoreTypes();
      if (core_types.empty())
        return AllCores();
  
      cpu_set_t set;
      CPU_ZERO(&set);
      for (size_t core_index = 0; core_index < core_types.size(); core_index++) {
        switch (core_types[core_index]) {
          case CPU::CoreType::kUnknown:
          case CPU::CoreType::kOther:
          case CPU::CoreType::kSymmetric:
            // In the presence of an unknown core type or symmetric architecture,
            // fall back to allowing all cores.
            return AllCores();
          case CPU::CoreType::kBigLittle_Little:
          case CPU::CoreType::kBigLittleBigger_Little:
            CPU_SET(core_index, &set);
            break;
          case CPU::CoreType::kBigLittle_Big:
          case CPU::CoreType::kBigLittleBigger_Big:
          case CPU::CoreType::kBigLittleBigger_Bigger:
            break;
        }
      }
      return set;
    }();
    return kLittleCores;
  }
  
  const cpu_set_t& BigCores() {
    static const cpu_set_t kBigCores = []() {
      const std::vector<CPU::CoreType>& core_types = CPU::GetGuessedCoreTypes();
      if (core_types.empty())
        return AllCores();
  
      cpu_set_t set;
      CPU_ZERO(&set);
      for (size_t core_index = 0; core_index < core_types.size(); core_index++) {
        switch (core_types[core_index]) {
          case CPU::CoreType::kUnknown:
          case CPU::CoreType::kOther:
          case CPU::CoreType::kSymmetric:
            // In the presence of an unknown core type or symmetric architecture,
            // fall back to allowing all cores.
            return AllCores();
          case CPU::CoreType::kBigLittle_Little:
          case CPU::CoreType::kBigLittleBigger_Little:
            break;
          case CPU::CoreType::kBigLittle_Big:
          case CPU::CoreType::kBigLittleBigger_Big:
          case CPU::CoreType::kBigLittleBigger_Bigger:
            CPU_SET(core_index, &set);
            break;
        }
      }
      return set;
    }();
    return kBigCores;
  }
  
  const cpu_set_t& BiggerCores() {
    static const cpu_set_t kBiggerCores = []() {
      const std::vector<CPU::CoreType>& core_types = CPU::GetGuessedCoreTypes();
      if (core_types.empty())
        return AllCores();
  
      cpu_set_t set;
      CPU_ZERO(&set);
      for (size_t core_index = 0; core_index < core_types.size(); core_index++) {
        switch (core_types[core_index]) {
          case CPU::CoreType::kUnknown:
          case CPU::CoreType::kOther:
          case CPU::CoreType::kSymmetric:
            // In the presence of an unknown core type or symmetric architecture,
            // fall back to allowing all cores.
            return AllCores();
          case CPU::CoreType::kBigLittle_Little:
          case CPU::CoreType::kBigLittleBigger_Little:
          case CPU::CoreType::kBigLittle_Big:
          case CPU::CoreType::kBigLittleBigger_Big:
            break;
          case CPU::CoreType::kBigLittleBigger_Bigger:
            CPU_SET(core_index, &set);
            break;
        }
      }
      return set;
    }();
    return kBiggerCores;
  }
  
  }  // anonymous namespace
  
  bool HasBigCpuCores() {
    static const bool kHasBigCores = []() {
      const std::vector<CPU::CoreType>& core_types = CPU::GetGuessedCoreTypes();
      if (core_types.empty())
        return false;
      for (CPU::CoreType core_type : core_types) {
        switch (core_type) {
          case CPU::CoreType::kUnknown:
          case CPU::CoreType::kOther:
          case CPU::CoreType::kSymmetric:
            return false;
          case CPU::CoreType::kBigLittle_Little:
          case CPU::CoreType::kBigLittleBigger_Little:
          case CPU::CoreType::kBigLittle_Big:
          case CPU::CoreType::kBigLittleBigger_Big:
          case CPU::CoreType::kBigLittleBigger_Bigger:
            return true;
        }
      }
      return false;
    }();
    return kHasBigCores;
  }
  
  bool HasBiggerCpuCores() {
    static const bool kHasBiggerCores = []() {
      const std::vector<CPU::CoreType>& core_types = CPU::GetGuessedCoreTypes();
      if (core_types.empty())
        return false;
      for (CPU::CoreType core_type : core_types) {
        switch (core_type) {
          case CPU::CoreType::kUnknown:
          case CPU::CoreType::kOther:
          case CPU::CoreType::kSymmetric:
            return false;
          case CPU::CoreType::kBigLittle_Little:
          case CPU::CoreType::kBigLittleBigger_Little:
          case CPU::CoreType::kBigLittle_Big:
          case CPU::CoreType::kBigLittleBigger_Big:
            continue;
          case CPU::CoreType::kBigLittleBigger_Bigger:
            return true;
        }
      }
      return false;
    }();
    return kHasBiggerCores;
  }
  
  bool SetThreadCpuAffinityMode(PlatformThreadId thread_id,
                                CpuAffinityMode affinity) {
    int result = 0;
    switch (affinity) {
      case CpuAffinityMode::kDefault: {
        const cpu_set_t& all_cores = AllCores();
        result = sched_setaffinity(thread_id, sizeof(all_cores), &all_cores);
        break;
      }
      case CpuAffinityMode::kLittleCoresOnly: {
        const cpu_set_t& little_cores = LittleCores();
        result =
            sched_setaffinity(thread_id, sizeof(little_cores), &little_cores);
        break;
      }
      case CpuAffinityMode::kBigCoresOnly: {
        const cpu_set_t& big_cores = BigCores();
        result = sched_setaffinity(thread_id, sizeof(big_cores), &big_cores);
        break;
      }
      case CpuAffinityMode::kBiggerCoresOnly: {
        const cpu_set_t& bigger_cores = BiggerCores();
        result =
            sched_setaffinity(thread_id, sizeof(bigger_cores), &bigger_cores);
        break;
      }
    }
    return result == 0;
  }
  
  bool SetProcessCpuAffinityMode(ProcessHandle process_handle,
                                 CpuAffinityMode affinity) {
    bool any_threads = false;
    bool result = true;
  
    internal::ForEachProcessTask(
        process_handle, [&any_threads, &result, affinity](
                            PlatformThreadId tid, const FilePath& /*task_path*/) {
          any_threads = true;
          result &= SetThreadCpuAffinityMode(tid, affinity);
        });
  
    return any_threads && result;
  }
  
  absl::optional<CpuAffinityMode> CurrentThreadCpuAffinityMode() {
    if (HasBigCpuCores()) {
      cpu_set_t set;
      sched_getaffinity(PlatformThread::CurrentId(), sizeof(set), &set);
      if (CPU_EQUAL(&set, &AllCores()))
        return CpuAffinityMode::kDefault;
      if (CPU_EQUAL(&set, &LittleCores()))
        return CpuAffinityMode::kLittleCoresOnly;
      if (CPU_EQUAL(&set, &BigCores()))
        return CpuAffinityMode::kBigCoresOnly;
      if (CPU_EQUAL(&set, &BiggerCores()))
        return CpuAffinityMode::kBiggerCoresOnly;
    }
    return absl::nullopt;
  }
  
  }  // namespace base
  