Make context switch for deserialized traces.

!!!WARNING!!! !!!WARNING!!! !!!WARNING!!!

This change CHANGES API but this change IS NOT ENFORCED statically.
After this change all accesses to deserialized trace points MUST be wrapped into LazyTraceReader.inContext() or withTraceContext() calls or preceded by switchGlobalContext() call.

!!!WARNING!!! !!!WARNING!!! !!!WARNING!!!

Author: lev-serebryakov-jetbrains

common/src/main/org/jetbrains/lincheck/trace/TraceContext.kt

@@ -19,7 +19,32 @@ import org.jetbrains.lincheck.descriptors.MethodSignature
 import org.jetbrains.lincheck.descriptors.VariableDescriptor
 import org.jetbrains.lincheck.descriptors.Types
 
-val TRACE_CONTEXT: TraceContext = TraceContext()
+var TRACE_CONTEXT: TraceContext = TraceContext()
+    private set
+
+val TRACE_CONTEXT_LOCK = Any()
+
+fun setGlobalTraceContext(traceContext: TraceContext) {
+    TRACE_CONTEXT = traceContext
+}
+
+/**
+ * Run [block] with global trace context set to [context], restore global after run.
+ *
+ * This function is effectively `synchronized`, it prevents concurrent access
+ * to global trace context to avoid races.
+ */
+inline fun <R> withTraceContext(context: TraceContext, block: () -> R): R {
+    synchronized(TRACE_CONTEXT_LOCK) {
+        val saveContext = TRACE_CONTEXT
+        setGlobalTraceContext(context)
+        try {
+            return block()
+        } finally {
+            setGlobalTraceContext(saveContext)
+        }
+    }
+}
 
 const val UNKNOWN_CODE_LOCATION_ID = -1
 

trace/src/main/org/jetbrains/lincheck/trace/Deserialization.kt

@@ -193,8 +193,7 @@ class LazyTraceReader(
             tracePoint = readTracePointWithChildAddresses()
         )
 
-    // TODO: Create new
-    val context: TraceContext = TRACE_CONTEXT
+    val context: TraceContext = TraceContext()
 
     private val dataStream: SeekableInputStream
     private val data: SeekableDataInput
@@ -220,81 +219,125 @@ class LazyTraceReader(
         index?.close()
     }
 
+    /**
+     * Run block with trace context belonging to this reader and restore global context afterward.
+     *
+     * @See withTraceContext
+     */
+    inline fun <R> inContext(block: () -> R): R = withTraceContext(context, block)
+
+    /**
+     * Switch global context to one belonging to this reader.
+     *
+     * This function is inherently unsafe: it doesn't prevent future switches of global context
+     * by other code or threads.
+     *
+     * Please, use [inContext] instead if possible.
+     */
+    @Deprecated("Use inContext if possible.")
+    fun switchGlobalContext() = setGlobalTraceContext(context)
+
+    /**
+     * Read top-level call tracepoints for each thread.
+     *
+     * This method runs in context of this reader's context and doesn't require [inContext].
+     */
     fun readRoots(): List<TRTracePoint> {
-        var start = System.currentTimeMillis()
-        loadContext()
-        Logger.debug { "Context loaded in ${System.currentTimeMillis() - start} ms" }
-        start = System.currentTimeMillis()
-
-        val roots = mutableMapOf<Int, TRTracePoint>()
-
-        dataBlocks.forEach {
-            val (threadId, blocks) = it
-            data.seek(blocks.first().physicalStart)
-            val kind = data.readKind()
-            check(kind == ObjectKind.BLOCK_START) { "Thread $threadId block 0 has wrong start: $kind" }
-            val blockId = data.readInt()
-            check(blockId == threadId) { "Thread $threadId block 0 has wrong idt: $blockId" }
-
-            val tracepoints = mutableListOf<TRTracePoint>()
-            loadTracePoints(
-                threadId = threadId,
-                maxRead = Integer.MAX_VALUE,
-                reader = this::readTracePointWithPostprocessor,
-                registrator = { _, tracePoint, _ ->
-                    tracepoints.add(tracePoint)
-                }
-            )
-            if (tracepoints.isEmpty()) {
-                System.err.println("Thread $threadId doesn't write any tracepoints")
-            } else {
-                if (tracepoints.size > 1) {
-                    System.err.println("Thread $threadId wrote too many root tracepoints: ${tracepoints.size}")
+        return withTraceContext(context) {
+            var start = System.currentTimeMillis()
+            loadContext()
+            Logger.debug { "Context loaded in ${System.currentTimeMillis() - start} ms" }
+            start = System.currentTimeMillis()
+
+            val roots = mutableMapOf<Int, TRTracePoint>()
+
+            dataBlocks.forEach {
+                val (threadId, blocks) = it
+                data.seek(blocks.first().physicalStart)
+                val kind = data.readKind()
+                check(kind == ObjectKind.BLOCK_START) { "Thread $threadId block 0 has wrong start: $kind" }
+                val blockId = data.readInt()
+                check(blockId == threadId) { "Thread $threadId block 0 has wrong idt: $blockId" }
+
+                val tracepoints = mutableListOf<TRTracePoint>()
+                loadTracePoints(
+                    threadId = threadId,
+                    maxRead = Integer.MAX_VALUE,
+                    reader = this::readTracePointWithPostprocessor,
+                    registrator = { _, tracePoint, _ ->
+                        tracepoints.add(tracePoint)
+                    }
+                )
+                if (tracepoints.isEmpty()) {
+                    System.err.println("Thread $threadId doesn't write any tracepoints")
+                } else {
+                    if (tracepoints.size > 1) {
+                        System.err.println("Thread $threadId wrote too many root tracepoints: ${tracepoints.size}")
+                    }
+                    roots[threadId] = tracepoints.first()
                 }
-                roots[threadId] = tracepoints.first()
             }
-        }
-        Logger.debug { "Roots loaded in ${System.currentTimeMillis() - start} ms" }
+            Logger.debug { "Roots loaded in ${System.currentTimeMillis() - start} ms" }
 
-        return roots.entries.sortedBy { it.key }.map { (_, tracePoint) -> tracePoint }
+            roots.entries.sortedBy { it.key }.map { (_, tracePoint) -> tracePoint }
+        }
     }
 
+    /**
+     * Load one level of children of giver call tracepoint.
+     *
+     * This method runs in context of this reader's context and doesn't require [inContext].
+     */
     fun loadAllChildren(parent: TRMethodCallTracePoint) {
-        val (start, end) = callTracepointChildren[parent.eventId]
-            ?: error("TRMethodCallTracePoint ${parent.eventId} is not found in index")
+        withTraceContext(context) {
+            val (start, end) = callTracepointChildren[parent.eventId]
+                ?: error("TRMethodCallTracePoint ${parent.eventId} is not found in index")
 
-        data.seek(calculatePhysicalOffset(parent.threadId, start))
+            data.seek(calculatePhysicalOffset(parent.threadId, start))
 
-        loadTracePoints(
-            threadId = parent.threadId,
-            maxRead = Integer.MAX_VALUE,
-            reader = this::readTracePointWithPostprocessor,
-            registrator = { idx, tracePoint, _ ->
-                parent.loadChild(idx, tracePoint)
-            }
-        )
+            loadTracePoints(
+                threadId = parent.threadId,
+                maxRead = Integer.MAX_VALUE,
+                reader = this::readTracePointWithPostprocessor,
+                registrator = { idx, tracePoint, _ ->
+                    parent.loadChild(idx, tracePoint)
+                }
+            )
 
-        val actualFooterPos = data.position() - 1 // 1 is size of object kind
-        check(actualFooterPos == calculatePhysicalOffset(parent.threadId, end)) {
-            "Input contains broken data: expected Tracepoint Footer for event ${parent.eventId} at position $end, got $actualFooterPos"
+            val actualFooterPos = data.position() - 1 // 1 is size of object kind
+            check(actualFooterPos == calculatePhysicalOffset(parent.threadId, end)) {
+                "Input contains broken data: expected Tracepoint Footer for event ${parent.eventId} at position $end, got $actualFooterPos"
+            }
         }
     }
 
+    /**
+     * Load child of giver call tracepoint at given index.
+     *
+     * This method runs in context of this reader's context and doesn't require [inContext].
+     */
     fun loadChild(parent: TRMethodCallTracePoint, childIdx: Int): Unit = loadChildrenRange(parent, childIdx, 1)
 
+    /**
+     * Load range of children of giver call tracepoint at given span.
+     *
+     * This method runs in context of this reader's context and doesn't require [inContext].
+     */
     fun loadChildrenRange(parent: TRMethodCallTracePoint, from: Int, count: Int) {
-        require(from in 0 ..< parent.events.size) { "From index $from must be in range 0..<${parent.events.size}" }
-        require(count in 1 .. parent.events.size - from) { "Count $count must be in range 1..${parent.events.size - from}" }
+        withTraceContext(context) {
+            require(from in 0..<parent.events.size) { "From index $from must be in range 0..<${parent.events.size}" }
+            require(count in 1..parent.events.size - from) { "Count $count must be in range 1..${parent.events.size - from}" }
 
-        data.seek(calculatePhysicalOffset(parent.threadId, parent.getChildAddress(from)))
-        loadTracePoints(
-            threadId = parent.threadId,
-            maxRead = count,
-            reader = this::readTracePointWithPostprocessor,
-            registrator = { idx, tracePoint, _ ->
-                parent.loadChild(idx + from, tracePoint)
-            }
-        )
+            data.seek(calculatePhysicalOffset(parent.threadId, parent.getChildAddress(from)))
+            loadTracePoints(
+                threadId = parent.threadId,
+                maxRead = count,
+                reader = this::readTracePointWithPostprocessor,
+                registrator = { idx, tracePoint, _ ->
+                    parent.loadChild(idx + from, tracePoint)
+                }
+            )
+        }
     }
 
     fun getChildAndRestorePosition(parent: TRMethodCallTracePoint, childIdx: Int): TRTracePoint? {
@@ -436,7 +479,6 @@ class LazyTraceReader(
     }
 
     private fun loadContextWithoutIndex() {
-        val context = TRACE_CONTEXT
         // Two Longs is header
         data.seek((Long.SIZE_BYTES * 2).toLong())
         loadAllObjectsDeep(
@@ -570,14 +612,31 @@ class LazyTraceReader(
 data class TraceWithContext(
     val context: TraceContext,
     val roots: List<TRTracePoint>
-)
+) {
+    /**
+     * Run block with trace context belonging to this trace and restore global context afterward.
+     *
+     * @See withTraceContext
+     */
+    inline fun <R> inContext(block: () -> R): R = withTraceContext(context, block)
+
+    /**
+     * Switch global context to one belonging to this trace.
+     *
+     * This function is inherently unsafe: it doesn't prevent future switches of global context
+     * by other code or threads.
+     *
+     * Please, use [inContext] instead if possible.
+     */
+    @Deprecated("Use inContext if possible.")
+    fun switchGlobalContext() = setGlobalTraceContext(context)
+}
 
 fun loadRecordedTrace(inp: InputStream): TraceWithContext {
     DataInputStream(inp.buffered(INPUT_BUFFER_SIZE)).use { input ->
         checkDataHeader(input)
 
-        // TODO: Create empty fresh context
-        val context = TRACE_CONTEXT
+        val context = TraceContext()
         val roots = mutableMapOf<Int, MutableList<TRTracePoint>>()
 
         loadAllObjectsDeep(

trace/src/main/org/jetbrains/lincheck/trace/Printing.kt

@@ -24,12 +24,13 @@ fun printRecorderTrace(fileName: String?, context: TraceContext, rootCallsPerThr
     )
 
 fun printRecorderTrace(output: OutputStream, context: TraceContext, rootCallsPerThread: List<TRTracePoint>, verbose: Boolean) {
-    check(context == TRACE_CONTEXT) { "Now only global TRACE_CONTEXT is supported" }
-    PrintStream(output.buffered(OUTPUT_BUFFER_SIZE)).use { output ->
-        val appendable = DefaultTRTextAppendable(output, verbose)
-        rootCallsPerThread.forEachIndexed { i, root ->
-            output.println("# Thread ${i+1}")
-            printTRPoint(appendable, root, 0)
+    withTraceContext(context) {
+        PrintStream(output.buffered(OUTPUT_BUFFER_SIZE)).use { output ->
+            val appendable = DefaultTRTextAppendable(output, verbose)
+            rootCallsPerThread.forEachIndexed { i, root ->
+                output.println("# Thread ${i + 1}")
+                printTRPoint(appendable, root, 0)
+            }
         }
     }
 }