// Copyright 2019 The TCMalloc 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
//
//     https://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.
#include "tcmalloc/internal/affinity.h"

#include <fcntl.h>
#include <stdarg.h>
#include <unistd.h>

#include <vector>

#include "absl/log/log.h"
#include "absl/types/span.h"
#include "tcmalloc/internal/config.h"
#include "tcmalloc/internal/cpu_utils.h"
#include "tcmalloc/internal/logging.h"

GOOGLE_MALLOC_SECTION_BEGIN
namespace tcmalloc {
namespace tcmalloc_internal {

std::vector<int> AllowedCpus() {
  // We have no need for dynamically sized sets (currently >1024 CPUs for glibc)
  // at the present time.  We could change this in the future.
  CpuSet allowed_cpus;
  TC_CHECK(allowed_cpus.GetAffinity(0));
  int n = allowed_cpus.Count(), c = 0;

  std::vector<int> result(n);
  for (int i = 0; i < kMaxCpus && n; i++) {
    if (allowed_cpus.IsSet(i)) {
      result[c++] = i;
      n--;
    }
  }
  TC_CHECK_EQ(0, n);

  return result;
}

static CpuSet SpanToCpuSetT(absl::Span<int> mask) {
  CpuSet result;
  result.Zero();
  for (int cpu : mask) {
    result.Set(cpu);
  }
  return result;
}

ScopedAffinityMask::ScopedAffinityMask(absl::Span<int> allowed_cpus) {
  specified_cpus_ = SpanToCpuSetT(allowed_cpus);
  // getaffinity should never fail.
  TC_CHECK(original_cpus_.GetAffinity(0));
  // See destructor comments on setaffinity interactions.  Tampered() will
  // necessarily be true in this case, especially when SetAffinity(0)
  // fails.
  if (!specified_cpus_.SetAffinity(0)) {
    PLOG(ERROR) << "SetAffinity(0) failed! ";
  }
}

ScopedAffinityMask::ScopedAffinityMask(int allowed_cpu) {
  specified_cpus_.Zero();
  specified_cpus_.Set(allowed_cpu);

  // getaffinity should never fail.
  TC_CHECK(original_cpus_.GetAffinity(0));
  // See destructor comments on setaffinity interactions.  Tampered() will
  // necessarily be true in this case, especially when SetAffinity(0)
  // fails.
  if (!specified_cpus_.SetAffinity(0)) {
    PLOG(ERROR) << "SetAffinity(0) failed! ";
  }
}

ScopedAffinityMask::~ScopedAffinityMask() {
  // If something else has already reset our affinity, do not attempt to
  // restrict towards our original mask.  This is best-effort as the tampering
  // may obviously occur during the destruction of *this.
  if (!Tampered()) {
    // Note:  We do not assert success here, conflicts may restrict us from all
    // 'original_cpus_'.
    (void)original_cpus_.SetAffinity(0);
  }
}

bool ScopedAffinityMask::Tampered() {
  CpuSet current_cpus;
  TC_CHECK(current_cpus.GetAffinity(0));
  return !CPU_EQUAL_S(kCpuSetBytes, &current_cpus,
                      &specified_cpus_);  // Mismatch => modified.
}

}  // namespace tcmalloc_internal
}  // namespace tcmalloc
GOOGLE_MALLOC_SECTION_END
