tsan: fix a bug in trace part switching

[dvyukov]

Summary

Callers of TraceSwitchPart expect that TraceAcquire will always succeed
after the call. It's possible that TryTraceFunc/TraceMutexLock in TraceSwitchPart
that restore the current stack/mutexset filled the trace part exactly up
to the TracePart::kAlignment gap and the next TraceAcquire won't succeed.
Skip the alignment gap after writing initial stack/mutexset to avoid that.

compiler-rt/lib/tsan/rtl/tsan_rtl.cpp

TraceTime(thr);
     }
   }
   for (uptr i = 0; i < thr->mset.Size(); i++) {
     MutexSet::Desc d = thr->mset.Get(i);
     for (uptr i = 0; i < d.count; i++)
       TraceMutexLock(thr, d.write ? EventType::kLock : EventType::kRLock, 0,
                      d.addr, d.stack_id);
   }
+  // Callers of TraceSwitchPart expect that TraceAcquire will always succeed
+  // after the call. It's possible that TryTraceFunc/TraceMutexLock above
+  // filled the trace part exactly up to the TracePart::kAlignment gap
+  // and the next TraceAcquire won't succeed. Skip the gap to avoid that.
+  EventFunc *ev;
+  if (!TraceAcquire(thr, &ev)) {
+    CHECK(TraceSkipGap(thr));
+    CHECK(TraceAcquire(thr, &ev));
+  }
   {
     Lock lock(&ctx->slot_mtx);
     // There is a small chance that the slot may be not queued at this point.
     // This can happen if the slot has kEpochLast epoch and another thread
     // in FindSlotAndLock discovered that it's exhausted and removed it from
     // the slot queue. kEpochLast can happen in 2 cases: (1) if TraceSwitchPart
     // was called with the slot locked and epoch already at kEpochLast,
     // or (2) if we've acquired a new slot in SlotLock in the beginning

compiler-rt/lib/tsan/tests/unit/tsan_trace_test.cpp

TRACE_TEST(Trace, MultiPart) {
   CHECK_EQ(mset.Get(0).addr, 0x5000);
   CHECK_EQ(mset.Get(0).write, true);
   CHECK_EQ(mset.Get(0).count, 2);
   CHECK_EQ(mset.Get(1).addr, 0x5001);
   CHECK_EQ(mset.Get(1).write, false);
   CHECK_EQ(mset.Get(1).count, 1);
 }
 
+TRACE_TEST(Trace, DeepSwitch) {
+  ThreadArray<1> thr;
+  for (int i = 0; i < 2000; i++) {
+    FuncEntry(thr, 0x1000);
+    const uptr kEvents = sizeof(TracePart) / sizeof(Event);
+    for (uptr i = 0; i < kEvents; i++) {
+      TraceMutexLock(thr, EventType::kLock, 0x4000, 0x5000, 0x6000);
+      TraceMutexUnlock(thr, 0x5000);
+    }
+  }
+}
+
 void CheckTraceState(uptr count, uptr finished, uptr excess, uptr recycle) {
   Lock l(&ctx->slot_mtx);
   Printf("CheckTraceState(%zu/%zu, %zu/%zu, %zu/%zu, %zu/%zu)\n",
          ctx->trace_part_total_allocated, count,
          ctx->trace_part_recycle_finished, finished,
          ctx->trace_part_finished_excess, excess,
          ctx->trace_part_recycle.Size(), recycle);
   CHECK_EQ(ctx->trace_part_total_allocated, count);