From f5033d1e88f4a9632611f0f2aeffbaaea774a2d4 Mon Sep 17 00:00:00 2001 From: Ian Lynagh Date: Fri, 4 Oct 2024 16:11:26 +0100 Subject: [PATCH] KE2: Make the shared stuff threadsafe --- .../src/main/kotlin/KotlinUsesExtractor.kt | 2 +- .../src/main/kotlin/TrapWriter.kt | 60 +++++++++++++++---- .../src/main/kotlin/utils/Logger.kt | 8 ++- 3 files changed, 53 insertions(+), 17 deletions(-) diff --git a/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt b/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt index 5a4576dc953..41a5285b499 100644 --- a/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt +++ b/java/kotlin-extractor2/src/main/kotlin/KotlinUsesExtractor.kt @@ -96,7 +96,7 @@ open class KotlinUsesExtractor( val pkg = f.packageFqName.asString() val jvmName = getFileClassName(f) val id = extractFileClass(pkg, jvmName) - if (tw.lm.fileClassLocationsExtracted.add(f)) { + if (tw.lm.markFileClassLocationAsExtracted(f)) { val fileId = tw.mkFileId(f.virtualFilePath, false) val locId = tw.getWholeFileLocation(fileId) tw.writeHasLocation(id, locId) diff --git a/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt b/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt index b6e06d05540..595e7ba3d0f 100644 --- a/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt +++ b/java/kotlin-extractor2/src/main/kotlin/TrapWriter.kt @@ -9,6 +9,8 @@ import com.semmle.extractor.java.PopulateFile import com.semmle.util.unicode.UTF8Util import java.io.BufferedWriter import java.io.File +import java.util.concurrent.locks.ReentrantLock +import kotlin.concurrent.withLock /* OLD: KE1 import org.jetbrains.kotlin.ir.IrElement @@ -28,16 +30,34 @@ import org.jetbrains.kotlin.psi.* * names, and maintains a mapping from keys (`@"..."`) to labels. */ class TrapLabelManager { + /** + * The lock that controls access to the label manager state. + * While we can make a thread-safe MutableMap with + * Collections.synchronizedMap and use getOrPut, that doesn't + * guarantee not to run the `defaultValue` function when it isn't + * necessary, which makes it useless for our purposes. + * TODO: We only actually need this for the diagnostic TRAP file. Make it optional? + */ + private val lock = ReentrantLock() + /** The next integer to use as a label name. */ private var nextInt: Int = 100 /** Returns a fresh label. */ fun getFreshLabel(): Label { - return IntLabel(nextInt++) + lock.withLock { + return IntLabel(nextInt++) + } } /** A mapping from a key (`@"..."`) to the label defined to be that key, if any. */ - val labelMapping: MutableMap> = mutableMapOf>() + private val labelMapping: MutableMap> = mutableMapOf>() + + fun withLabelMapping(action: (MutableMap>) -> T): T { + lock.withLock { + return action(labelMapping) + } + } /* OLD: KE1 @@ -60,13 +80,25 @@ class TrapLabelManager { * This allows us to keep track of whether we've written the location already in this TRAP file, * to avoid duplication. */ - val fileClassLocationsExtracted = HashSet() + private val fileClassLocationsExtracted = HashSet() + + /** + * Indicate that we want `file`'s file class location marked as extracted. + * Returns true if we need to actually write the TRAP for it, or false + * if it's already been done. + */ + fun markFileClassLocationAsExtracted(file: KtFile): Boolean { + lock.withLock { + return fileClassLocationsExtracted.add(file) + } + } } /** * A `TrapWriter` is used to write TRAP to a particular TRAP file. There may be multiple * `TrapWriter`s for the same file, as different instances will have different additional state, but * they must all share the same `TrapLabelManager` and `BufferedWriter`. + * `BasicLogger`s, `TrapLabelManager` and `BufferedWriter` are threadsafe, so `TrapWriter`s are too. */ abstract class TrapWriter( protected val basicLogger: BasicLogger, @@ -85,7 +117,7 @@ abstract class TrapWriter( TODO: Inline this if it can remain private */ private fun getExistingLabelFor(key: String): Label? { - return lm.labelMapping.get(key)?.cast() + return lm.withLabelMapping { labelMapping -> labelMapping.get(key)?.cast() } } /** @@ -94,15 +126,17 @@ abstract class TrapWriter( */ @JvmOverloads // Needed so Java can call a method with an optional argument fun getLabelFor(key: String, initialise: (Label) -> Unit = {}): Label { - val maybeLabel: Label? = getExistingLabelFor(key) - if (maybeLabel == null) { - val label: Label = lm.getFreshLabel() - lm.labelMapping.put(key, label) - writeTrap("$label = $key\n") - initialise(label) - return label - } else { - return maybeLabel + return lm.withLabelMapping { labelMapping -> + val maybeLabel: Label? = getExistingLabelFor(key) + if (maybeLabel == null) { + val label: Label = lm.getFreshLabel() + labelMapping.put(key, label) + writeTrap("$label = $key\n") + initialise(label) + label + } else { + maybeLabel + } } } diff --git a/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt b/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt index e91ca9e6d9d..15cab0f1412 100644 --- a/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt +++ b/java/kotlin-extractor2/src/main/kotlin/utils/Logger.kt @@ -1,6 +1,7 @@ package com.github.codeql import com.intellij.psi.PsiElement +import java.io.BufferedWriter import java.io.File import java.io.FileWriter import java.io.OutputStreamWriter @@ -127,15 +128,16 @@ class LoggerBase(val diagnosticCounter: DiagnosticCounter) : BasicLogger { verbosity = System.getenv("CODEQL_EXTRACTOR_KOTLIN_VERBOSITY")?.toIntOrNull() ?: 3 } - private val logStream: Writer + // Use BufferedWriter as it is threadsafe + private val logStream: BufferedWriter init { val extractorLogDir = System.getenv("CODEQL_EXTRACTOR_JAVA_LOG_DIR") if (extractorLogDir == null || extractorLogDir == "") { - logStream = OutputStreamWriter(System.out) + logStream = BufferedWriter(OutputStreamWriter(System.out)) } else { val logFile = File.createTempFile("kotlin-extractor.", ".log", File(extractorLogDir)) - logStream = FileWriter(logFile) + logStream = BufferedWriter(FileWriter(logFile)) } }