[ml-local-models] handle invalid storage and limit number of classes info in memory

(cherry picked from commit c05e73559ee54d7c8bd90fb325f8e07a6b98e989)

IJ-MR-7427

GitOrigin-RevId: 1d821a314ee9d3b4cb5f00628989aa0dc7df6414
This commit is contained in:
Alexey Kalina
2021-03-29 12:25:11 +03:00
committed by intellij-monorepo-bot
parent dc2e58acf3
commit faecec4099
6 changed files with 43 additions and 13 deletions

View File

@@ -7,4 +7,13 @@ interface LocalModelBuilder {
fun onFinished()
fun fileVisitor(): PsiElementVisitor
fun build(): LocalModel?
companion object {
val DUMB_BUILDER = object : LocalModelBuilder {
override fun onStarted() = Unit
override fun onFinished() = Unit
override fun fileVisitor(): PsiElementVisitor = PsiElementVisitor.EMPTY_VISITOR
override fun build(): LocalModel? = null
}
}
}

View File

@@ -15,7 +15,7 @@ abstract class ClassesFrequencyModelFactory : FrequencyModelFactory<ClassesUsage
override fun modelBuilder(project: Project, language: Language): LocalModelBuilder {
val storagesPath = StorageUtil.storagePath(project, language)
val storage = ClassesFrequencyStorage.getStorage(storagesPath)
val storage = ClassesFrequencyStorage.getStorage(storagesPath) ?: return LocalModelBuilder.DUMB_BUILDER
return object : LocalModelBuilder {

View File

@@ -14,12 +14,11 @@ class ClassesFrequencyStorage internal constructor(private val storageDirectory:
companion object {
private const val STORAGE_NAME = "classes_frequency"
private const val VERSION = 1
private const val CLASS_FREQUENCY_THRESHOLD = 5
private const val MAX_CLASSES_IN_MEMORY = 1500
fun getStorage(baseDirectory: Path): ClassesFrequencyStorage {
fun getStorage(baseDirectory: Path): ClassesFrequencyStorage? {
val storageDirectory = baseDirectory.resolve(STORAGE_NAME)
StorageUtil.prepareStorage(storageDirectory, VERSION)
return ClassesFrequencyStorage(storageDirectory)
return StorageUtil.getStorage(storageDirectory, VERSION) { path -> ClassesFrequencyStorage(path) }
}
}
private var isValid: Boolean = true
@@ -73,15 +72,23 @@ class ClassesFrequencyStorage internal constructor(private val storageDirectory:
memoryStorage.clear()
totalClasses = 0
totalClassesUsages = 0
val sortedClasses = sortedSetOf<Pair<String, Int>>(compareBy({ it.second }, { it.first }))
persistentStorage.processKeys(Processor {
val count = persistentStorage.get(it) ?: return@Processor true
totalClasses++
totalClassesUsages += count
if (count >= CLASS_FREQUENCY_THRESHOLD) {
memoryStorage[it] = count
if (totalClasses > MAX_CLASSES_IN_MEMORY) {
val first = sortedClasses.first()
if (first.second < count) {
sortedClasses.remove(first)
sortedClasses.add(Pair(it, count))
}
} else {
sortedClasses.add(Pair(it, count))
}
return@Processor true
})
sortedClasses.forEach { memoryStorage[it.first] = it.second }
persistentStorage.force()
}
}

View File

@@ -16,7 +16,7 @@ abstract class MethodsFrequencyModelFactory : FrequencyModelFactory<MethodsUsage
override fun modelBuilder(project: Project, language: Language): LocalModelBuilder {
val storagesPath = StorageUtil.storagePath(project, language)
val storage = MethodsFrequencyStorage.getStorage(storagesPath)
val storage = MethodsFrequencyStorage.getStorage(storagesPath) ?: return LocalModelBuilder.DUMB_BUILDER
return object : LocalModelBuilder {

View File

@@ -17,10 +17,9 @@ class MethodsFrequencyStorage internal constructor(private val storageDirectory:
private const val STORAGE_NAME = "methods_frequency"
private const val VERSION = 1
fun getStorage(baseDirectory: Path): MethodsFrequencyStorage {
fun getStorage(baseDirectory: Path): MethodsFrequencyStorage? {
val storageDirectory = baseDirectory.resolve(STORAGE_NAME)
StorageUtil.prepareStorage(storageDirectory, VERSION)
return MethodsFrequencyStorage(storageDirectory)
return StorageUtil.getStorage(storageDirectory, VERSION) { path -> MethodsFrequencyStorage(path) }
}
}
private var isValid: Boolean = true

View File

@@ -47,10 +47,25 @@ object StorageUtil {
}
}
fun prepareStorage(storageDirectory: Path, version: Int) {
fun <T> getStorage(storageDirectory: Path, version: Int, factory: (Path) -> T): T? {
try {
prepareStorage(storageDirectory, version)
return factory(storageDirectory)
} catch (t: Throwable) {
try {
prepareStorage(storageDirectory, version, forceDelete = true)
return factory(storageDirectory)
} catch (t: Throwable) {
LOG.error(t)
return null
}
}
}
private fun prepareStorage(storageDirectory: Path, version: Int, forceDelete: Boolean = false) {
if (storageDirectory.exists()) {
val info = readInfo(storageDirectory)
if (info == null || !info.isValid || info.version != version) {
if (forceDelete || info == null || !info.isValid || info.version != version) {
storageDirectory.delete()
}
}