[scudo][standalone] Add bytemap classes

[cryptoad]

Summary

The bytemap classes will be used by the primary32 allocator to associate
classes with memory regions. It's similar to the sanitizer_common one
except for the fact that the base (level1) maps are mapped instead of
being static to reduce the memory footprint of an uninitialized allocator.

[morehouse]

lib/scudo/standalone/bytemap.h

23

Having two init functions is a little confusing. Would it make more sense to have a constexpr constructor and a single init() function?

lib/scudo/standalone/tests/bytemap_test.cc

15

Tests for FlatByteMap?

[cryptoad]

Adding test for the flat bytemap, reorganzing the existing twolevel one.

[cryptoad]

lib/scudo/standalone/bytemap.h

23

With this I followed the way it was done in sanitizer_common with an init function that would skip the zero initialization if not needed (most our structures are either mapped or static). Would you have an example of what you are suggesting, I am not sure I fully understand how to organize it?

[morehouse]

lib/scudo/standalone/bytemap.h

23

Essentially, make the zero-initialization explicit in the constexpr constructor, and then do all dynamic initialization in init.

class FlatByteMap {
  constexpr FlatByteMap() : Map(nullptr) {}
  void init() {
    Map = reinterpret_cast<u8 *>(map(nullptr, Size, "scudo:bytemap"));
  }
};

class TwoLevelByteMap {
  constexpr TwoLevelByteMap() : Level1Map(nullptr), Mutex{} {}
  void init() {
    Level1Map = reinterpret_cast<atomic_uptr *>(
        map(nullptr, sizeof(atomic_uptr) * Level1Size, "scudo:bytemap"));
  }
};

The constexpr constructor will initialize any statics / globals at compile time while still working at runtime for non-static instantiations. Then the call to init completes any dynamic initialization required.

[morehouse]

lib/scudo/standalone/bytemap.h

23

The linker-init function will work fine, but since we're reimplementing much of sanitizer_common for Scudo now, maybe it's a good opportunity to modernize a bit.

[morehouse]

lib/scudo/standalone/bytemap.h

23

You'll probably also need to refactor StaticSpinMutex to use a constexpr constructor for initialization. In fact, we could probably completely remove StaticSpinMutex and just give SpinMutex a constexpr constructor.

class SpinMutex {
  constexpr SpinMutex() : State{} {}
};

Then SpinMutex can be used as either global or local.

[cryptoad]

lib/scudo/standalone/bytemap.h

23

I went into the rabbit hole of changing those and started hitting issues along the way (with __thread).
If that's Ok with you I'd like to land this in it's current shape and revisit that later.

[morehouse]

lib/scudo/standalone/bytemap.h

23

SGTM

lib/scudo/standalone/CMakeLists.txt

 
 set(SCUDO_SOURCES
   common.cc
   fuchsia.cc
   linux.cc)
 
 set(SCUDO_HEADERS
   atomic_helpers.h
+  bytemap.h
   internal_defs.h
   linux.h
   list.h
   mutex.h
   platform.h
   vector.h)
 
 if(COMPILER_RT_HAS_SCUDO_STANDALONE)

lib/scudo/standalone/bytemap.h

+//===-- bytemap.h -----------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef SCUDO_BYTEMAP_H_
+#define SCUDO_BYTEMAP_H_
+
+#include "atomic_helpers.h"
+#include "common.h"
+#include "mutex.h"
+
+namespace scudo {
+
+template <uptr Size> class FlatByteMap {
+public:
+  void initLinkerInitialized() {
+    Map = reinterpret_cast<u8 *>(map(nullptr, Size, "scudo:bytemap"));
+  }
+  void init() { initLinkerInitialized(); }
+
+  void set(uptr Index, u8 Value) {
+    DCHECK_LT(Index, Size);
+    DCHECK_EQ(0U, Map[Index]);
+    Map[Index] = Value;
+  }
+  u8 operator[](uptr Index) {
+    DCHECK_LT(Index, Size);
+    return Map[Index];
+  }
+
+private:
+  u8 *Map;
+};
+
+template <uptr Level1Size, uptr Level2Size> class TwoLevelByteMap {
+public:
+  void initLinkerInitialized() {
+    Level1Map = reinterpret_cast<atomic_uptr *>(
+        map(nullptr, sizeof(atomic_uptr) * Level1Size, "scudo:bytemap"));
+  }
+  void init() {
+    initLinkerInitialized();
+    Mutex.init();
+  }
+
+  void reset() {
+    for (uptr I = 0; I < Level1Size; I++) {
+      u8 *P = get(I);
+      if (!P)
+        continue;
+      unmap(P, Level2Size);
+    }
+    memset(Level1Map, 0, sizeof(atomic_uptr) * Level1Size);
+  }
+
+  uptr size() const { return Level1Size * Level2Size; }
+
+  void set(uptr Index, u8 Value) {
+    DCHECK_LT(Index, Level1Size * Level2Size);
+    u8 *Level2Map = getOrCreate(Index / Level2Size);
+    DCHECK_EQ(0U, Level2Map[Index % Level2Size]);
+    Level2Map[Index % Level2Size] = Value;
+  }
+
+  u8 operator[](uptr Index) const {
+    DCHECK_LT(Index, Level1Size * Level2Size);
+    u8 *Level2Map = get(Index / Level2Size);
+    if (!Level2Map)
+      return 0;
+    return Level2Map[Index % Level2Size];
+  }
+
+private:
+  u8 *get(uptr Index) const {
+    DCHECK_LT(Index, Level1Size);
+    return reinterpret_cast<u8 *>(
+        atomic_load(&Level1Map[Index], memory_order_acquire));
+  }
+
+  u8 *getOrCreate(uptr Index) {
+    u8 *Res = get(Index);
+    if (!Res) {
+      SpinMutexLock L(&Mutex);
+      if (!(Res = get(Index))) {
+        Res = reinterpret_cast<u8 *>(map(nullptr, Level2Size, "scudo:bytemap"));
+        atomic_store(&Level1Map[Index], reinterpret_cast<uptr>(Res),
+                     memory_order_release);
+      }
+    }
+    return Res;
+  }
+
+  atomic_uptr *Level1Map;
+  StaticSpinMutex Mutex;
+};
+
+} // namespace scudo
+
+#endif // SCUDO_BYTEMAP_H_

lib/scudo/standalone/tests/CMakeLists.txt

foreach(arch ${SCUDO_TEST_ARCH})
         CFLAGS ${SCUDO_UNITTEST_CFLAGS}
         LINK_FLAGS ${LINK_FLAGS})
     endforeach()
   endif()
 endmacro()
 
 set(SCUDO_UNIT_TEST_SOURCES
   atomic_test.cc
+  bytemap_test.cc
   list_test.cc
   map_test.cc
   mutex_test.cc
   vector_test.cc
   scudo_unit_test_main.cc)
 
 add_scudo_unittest(ScudoUnitTest
   SOURCES ${SCUDO_UNIT_TEST_SOURCES})

lib/scudo/standalone/tests/bytemap_test.cc

+//===-- bytemap_test.cc------------------------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "bytemap.h"
+
+#include "gtest/gtest.h"
+
+#include <string.h>
+
+template <typename T> void testMap(T &Map, scudo::uptr Size) {
+  Map.init();
+  for (scudo::uptr I = 0; I < Size; I += 7)
+    Map.set(I, (I % 100) + 1);
+  for (scudo::uptr J = 0; J < Size; J++) {
+    if (J % 7)
+      EXPECT_EQ(Map[J], 0);
+    else
+      EXPECT_EQ(Map[J], (J % 100) + 1);
+  }
+}
+
+TEST(ScudoByteMapTest, FlatByteMap) {
+  const scudo::uptr Size = 1U << 10;
+  scudo::FlatByteMap<Size> Map;
+  testMap(Map, Size);
+}
+
+TEST(ScudoByteMapTest, TwoLevelByteMap) {
+  const scudo::uptr Size1 = 1U << 6, Size2 = 1U << 12;
+  scudo::TwoLevelByteMap<Size1, Size2> Map;
+  testMap(Map, Size1 * Size2);
+  Map.reset();
+}
+
+using TestByteMap = scudo::TwoLevelByteMap<1U << 12, 1U << 13>;
+
+struct TestByteMapParam {
+  TestByteMap *Map;
+  scudo::uptr Shard;
+  scudo::uptr NumberOfShards;
+};
+
+void *populateByteMap(void *Param) {
+  TestByteMapParam *P = reinterpret_cast<TestByteMapParam *>(Param);
+  for (scudo::uptr I = P->Shard; I < P->Map->size(); I += P->NumberOfShards) {
+    scudo::u8 V = static_cast<scudo::u8>((I % 100) + 1);
+    P->Map->set(I, V);
+    EXPECT_EQ((*P->Map)[I], V);
+  }
+  return 0;
+}
+
+TEST(ScudoByteMapTest, ThreadedTwoLevelByteMap) {
+  TestByteMap Map;
+  Map.init();
+  static const scudo::uptr NumberOfThreads = 16U;
+  pthread_t T[NumberOfThreads];
+  TestByteMapParam P[NumberOfThreads];
+  for (scudo::uptr I = 0; I < NumberOfThreads; I++) {
+    P[I].Map = &Map;
+    P[I].Shard = I;
+    P[I].NumberOfShards = NumberOfThreads;
+    pthread_create(&T[I], 0, populateByteMap, &P[I]);
+  }
+  for (scudo::uptr I = 0; I < NumberOfThreads; I++)
+    pthread_join(T[I], 0);
+  Map.reset();
+}