Skip to content

Commit

Permalink
Add an implementation of NumCPUs() that is offline CPU aware (#188).
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 547256513
Change-Id: I44c42b154241bbbba21efa0cce1e8e4bc0f6a625
  • Loading branch information
@ckennelly
ckennelly authored and Copybara-Service committed Jul 11, 2023
1 parent 65c9cbd commit 5823a86
Show file tree
Hide file tree
Showing 8 changed files with 389 additions and 197 deletions.
49 changes: 39 additions & 10 deletions tcmalloc/internal/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -409,6 +409,7 @@ cc_library(
":environment",
":logging",
":percpu",
":sysinfo",
":util",
"@com_google_absl//absl/base",
"@com_google_absl//absl/base:core_headers",
Expand All @@ -418,16 +419,6 @@ cc_library(
],
)

cc_fuzz_test(
name = "numa_fuzz",
testonly = 1,
srcs = ["numa_fuzz.cc"],
copts = TCMALLOC_DEFAULT_COPTS,
deps = [
":numa",
],
)

cc_test(
name = "numa_test",
srcs = ["numa_test.cc"],
Expand Down Expand Up @@ -811,6 +802,44 @@ cc_test(
],
)

cc_library(
name = "sysinfo",
srcs = ["sysinfo.cc"],
hdrs = ["sysinfo.h"],
copts = TCMALLOC_DEFAULT_COPTS,
deps = [
":config",
":logging",
":util",
"@com_google_absl//absl/functional:function_ref",
"@com_google_absl//absl/strings",
],
)

cc_fuzz_test(
name = "sysinfo_fuzz",
testonly = 1,
srcs = ["sysinfo_fuzz.cc"],
copts = TCMALLOC_DEFAULT_COPTS,
deps = [
":sysinfo",
],
)

cc_test(
name = "sysinfo_test",
srcs = ["sysinfo_test.cc"],
copts = TCMALLOC_DEFAULT_COPTS,
deps = [
":logging",
":sysinfo",
"@com_google_absl//absl/base",
"@com_google_absl//absl/random",
"@com_google_absl//absl/strings",
"@com_google_googletest//:gtest_main",
],
)

# An empty rule to force libc malloc instead of TCMalloc.
cc_library(
name = "system_malloc",
Expand Down
71 changes: 1 addition & 70 deletions tcmalloc/internal/numa.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,11 +30,11 @@
#include "absl/base/attributes.h"
#include "absl/base/internal/sysinfo.h"
#include "absl/functional/function_ref.h"
#include "absl/strings/numbers.h"
#include "absl/strings/string_view.h"
#include "tcmalloc/internal/config.h"
#include "tcmalloc/internal/environment.h"
#include "tcmalloc/internal/logging.h"
#include "tcmalloc/internal/sysinfo.h"
#include "tcmalloc/internal/util.h"

GOOGLE_MALLOC_SECTION_BEGIN
Expand All @@ -53,75 +53,6 @@ int OpenSysfsCpulist(size_t node) {
return signal_safe_open(path, O_RDONLY | O_CLOEXEC);
}

namespace {
bool IsInBounds(int cpu) { return 0 <= cpu && cpu < CPU_SETSIZE; }
} // namespace

std::optional<cpu_set_t> ParseCpulist(
absl::FunctionRef<ssize_t(char*, size_t)> read) {
cpu_set_t set;
CPU_ZERO(&set);

std::array<char, 16> buf;
size_t carry_over = 0;
int cpu_from = -1;

while (true) {
const ssize_t rc = read(buf.data() + carry_over, buf.size() - carry_over);
if (ABSL_PREDICT_FALSE(rc < 0)) {
return std::nullopt;
}

const absl::string_view current(buf.data(), carry_over + rc);

// If we have no more data to parse & couldn't read any then we've reached
// the end of the input & are done.
if (current.empty() && rc == 0) {
break;
}
if (current == "\n" && rc == 0) {
break;
}

size_t consumed;
const size_t dash = current.find('-');
const size_t comma = current.find(',');
if (dash != absl::string_view::npos && dash < comma) {
if (!absl::SimpleAtoi(current.substr(0, dash), &cpu_from) ||
!IsInBounds(cpu_from)) {
return std::nullopt;
}
consumed = dash + 1;
} else if (comma != absl::string_view::npos || rc == 0) {
int cpu;
if (!absl::SimpleAtoi(current.substr(0, comma), &cpu) ||
!IsInBounds(cpu)) {
return std::nullopt;
}
if (comma == absl::string_view::npos) {
consumed = current.size();
} else {
consumed = comma + 1;
}
if (cpu_from != -1) {
for (int c = cpu_from; c <= cpu; c++) {
CPU_SET(c, &set);
}
cpu_from = -1;
} else {
CPU_SET(cpu, &set);
}
} else {
consumed = 0;
}

carry_over = current.size() - consumed;
memmove(buf.data(), buf.data() + consumed, carry_over);
}

return set;
}

bool InitNumaTopology(size_t cpu_to_scaled_partition[CPU_SETSIZE],
uint64_t* const partition_to_nodes,
NumaBindMode* const bind_mode,
Expand Down
14 changes: 0 additions & 14 deletions tcmalloc/internal/numa.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,20 +143,6 @@ class NumaTopology {
// returns the file descriptor.
int OpenSysfsCpulist(size_t node);

// Parse a CPU list in the format used by
// /sys/devices/system/node/nodeX/cpulist files - that is, individual CPU
// numbers or ranges in the format <start>-<end> inclusive all joined by comma
// characters.
//
// Returns absl::nullopt on error.
//
// The read function is expected to operate much like the read syscall. It
// should read up to `count` bytes into `buf` and return the number of bytes
// actually read. If an error occurs during reading it should return -1 with
// errno set to an appropriate error code.
std::optional<cpu_set_t> ParseCpulist(
absl::FunctionRef<ssize_t(char* buf, size_t count)> read);

// Initialize the data members of a NumaTopology<> instance.
//
// This function must only be called once per NumaTopology<> instance, and
Expand Down
102 changes: 0 additions & 102 deletions tcmalloc/internal/numa_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -241,108 +241,6 @@ TEST_F(NumaTopologyTest, Host) {
// we can do beyond checking that we didn't crash.
}

TEST(ParseCpulistTest, Empty) {
absl::string_view empty("\n");

const absl::optional<cpu_set_t> parsed =
ParseCpulist([&](char* const buf, const size_t count) -> ssize_t {
// Calculate how much data we have left to provide.
const size_t to_copy = std::min(count, empty.size());

// If none, we have no choice but to provide nothing.
if (to_copy == 0) return 0;

memcpy(buf, empty.data(), to_copy);
empty.remove_prefix(to_copy);
return to_copy;
});

// No CPUs should be active on this NUMA node.
ASSERT_THAT(parsed, testing::Ne(std::nullopt));
EXPECT_EQ(CPU_COUNT(&*parsed), 0);
}

TEST(ParseCpulistTest, NotInBounds) {
std::string cpulist = absl::StrCat("0-", CPU_SETSIZE);

const absl::optional<cpu_set_t> parsed =
ParseCpulist([&](char* const buf, const size_t count) -> ssize_t {
// Calculate how much data we have left to provide.
const size_t to_copy = std::min(count, cpulist.size());

// If none, we have no choice but to provide nothing.
if (to_copy == 0) return 0;

memcpy(buf, cpulist.data(), to_copy);
cpulist.erase(0, to_copy);
return to_copy;
});

ASSERT_THAT(parsed, testing::Eq(std::nullopt));
}

// Ensure that we can parse randomized cpulists correctly.
TEST(ParseCpulistTest, Random) {
absl::BitGen gen;

static constexpr int kIterations = 100;
for (int i = 0; i < kIterations; i++) {
cpu_set_t reference;
CPU_ZERO(&reference);

// Set a random number of CPUs within the reference set.
const double density = absl::Uniform(gen, 0.0, 1.0);
for (int cpu = 0; cpu < CPU_SETSIZE; cpu++) {
if (absl::Bernoulli(gen, density)) {
CPU_SET(cpu, &reference);
}
}

// Serialize the reference set into a cpulist-style string.
std::vector<std::string> components;
for (int cpu = 0; cpu < CPU_SETSIZE; cpu++) {
if (!CPU_ISSET(cpu, &reference)) continue;

const int start = cpu;
int next = cpu + 1;
while (next < CPU_SETSIZE && CPU_ISSET(next, &reference)) {
cpu = next;
next = cpu + 1;
}

if (cpu == start) {
components.push_back(absl::StrCat(cpu));
} else {
components.push_back(absl::StrCat(start, "-", cpu));
}
}
const std::string serialized = absl::StrJoin(components, ",");

// Now parse that string using our ParseCpulist function, randomizing the
// amount of data we provide to it from each read.
absl::string_view remaining(serialized);
const absl::optional<cpu_set_t> parsed =
ParseCpulist([&](char* const buf, const size_t count) -> ssize_t {
// Calculate how much data we have left to provide.
const size_t max = std::min(count, remaining.size());

// If none, we have no choice but to provide nothing.
if (max == 0) return 0;

// If we do have data, return a randomly sized subset of it to stress
// the logic around reading partial values.
const size_t copy = absl::Uniform(gen, static_cast<size_t>(1), max);
memcpy(buf, remaining.data(), copy);
remaining.remove_prefix(copy);
return copy;
});

// We ought to have parsed the same set of CPUs that we serialized.
ASSERT_THAT(parsed, testing::Ne(std::nullopt));
EXPECT_TRUE(CPU_EQUAL(&*parsed, &reference));
}
}

} // namespace
} // namespace tcmalloc_internal
} // namespace tcmalloc

0 comments on commit 5823a86

Please sign in to comment.