Impl Arena checkpoint.

and allow thread to use thread local scratch arena.

PiperOrigin-RevId: 680948868
Change-Id: Id740bc14007ce4fae8eb51ccdb1ab22f10b6ac15

src/tools/remote/src/main/cpp/testonly_output_service/memory.cc

@@ -42,6 +42,7 @@ Arena *AllocArena(size_t reserve_size) {
 }
 
 void FreeArena(Arena *arena) {
+  assert(arena->temp_memory_count == 0 && "Temporary memory still in use");
   size_t reserved_size = arena->reserved - (uint8_t *)arena;
   ReleaseMemory(arena, reserved_size);
 }
@@ -73,3 +74,50 @@ void PopArena(Arena *arena, size_t size) {
   arena->top -= size;
   assert(arena->top >= GetArenaBase(arena));
 }
+
+TemporaryMemory BeginTemporaryMemory(Arena *arena) {
+  TemporaryMemory temp;
+  temp.arena = arena;
+  temp.top = arena->top;
+  ++arena->temp_memory_count;
+  return temp;
+}
+
+void EndTemporaryMempory(TemporaryMemory temp) {
+  Arena *arena = temp.arena;
+  assert(arena->temp_memory_count > 0 && temp.top <= arena->top);
+  PopArena(arena, arena->top - temp.top);
+  assert(arena->top == temp.top);
+  --arena->temp_memory_count;
+}
+
+constexpr size_t kScratchArenaCount = 2;
+thread_local Arena *t_scratch_arenas[kScratchArenaCount];
+
+Arena *GetScratchArena(Arena **conflicts, size_t count) {
+  Arena *result = 0;
+  for (size_t i = 0; i < kScratchArenaCount; ++i) {
+    Arena *candidate = t_scratch_arenas[i];
+    if (!candidate) {
+      result = AllocArena();
+      t_scratch_arenas[i] = result;
+      break;
+    }
+
+    bool conflict = false;
+    for (size_t j = 0; j < count; ++j) {
+      if (candidate == conflicts[j]) {
+        conflict = true;
+        break;
+      }
+    }
+
+    if (!conflict) {
+      result = candidate;
+      break;
+    }
+  }
+
+  assert(result && "No scratch arena available");
+  return result;
+}

src/tools/remote/src/main/cpp/testonly_output_service/memory.h

@@ -50,9 +50,10 @@ struct Arena {
   //
   // Invariant: top <= committed <= reserved
   uint8_t *top;
+  uint32_t temp_memory_count;
 };
 
-Arena *AllocArena(size_t reserve_size);
+Arena *AllocArena(size_t reserve_size = GiB(1));
 void FreeArena(Arena *arena);
 
 // PushArena allocates a block of memory of the given size in the arena. The
@@ -63,4 +64,40 @@ void PopArena(Arena *arena, size_t size);
 #define PushArray(arena, type, count) \
   (type *)PushArena(arena, sizeof(type) * (count))
 
+struct TemporaryMemory {
+  Arena *arena;
+  uint8_t *top;
+};
+
+// Begins temporary memory allocations for the arena. The returned
+// TemporaryMemory must be passed to EndTemporaryMemory() when the memory is no
+// longer needed.
+TemporaryMemory BeginTemporaryMemory(Arena *arena);
+// Ends the temporary memory allocations. All memory allocated since
+// BeginTemporaryMemory() was called is freed.
+void EndTemporaryMempory(TemporaryMemory temp);
+
+// Gets a thread-local arena for temporary use. Use `conflicts` to avoid
+// returning the same arena as a previous call to GetScratchArena().
+Arena *GetScratchArena(Arena **conflicts, size_t count);
+
+// Begins the use of a thread-local scratch arena. The returned TemporaryMemory
+// must be passed to EndScratch() when the memory is no longer needed.
+static inline TemporaryMemory BeginScratch(Arena **conflicts, size_t count) {
+  return BeginTemporaryMemory(GetScratchArena(conflicts, count));
+}
+
+static inline TemporaryMemory BeginScratch(Arena *conflict) {
+  if (conflict) {
+    return BeginTemporaryMemory(GetScratchArena(&conflict, 1));
+  } else {
+    return BeginTemporaryMemory(GetScratchArena(0, 0));
+  }
+}
+
+// Ends the use of a thread-local scratch arena.
+static inline void EndScratch(TemporaryMemory scratch) {
+  EndTemporaryMempory(scratch);
+}
+
 #endif  // BAZEL_SRC_TOOLS_REMOTE_SRC_MAIN_CPP_TESTONLY_OUTPUT_SERVICE_MEMORY_H_