IDEA-383440 Identify data services entities and mark them

Change-Id: I32f0b99e90c492cb48629749e92ed3131eade926

Author: kerimsenol

plugins/gradle/generated/org/jetbrains/plugins/gradle/model/projectModel/impl/MetadataStorageImpl.kt

@@ -108,6 +108,7 @@
 
     <externalProjectDataService implementation="org.jetbrains.plugins.gradle.service.syncAction.impl.bridge.GradleBridgeProjectDataService"/>
     <externalWorkspaceDataService implementation="org.jetbrains.plugins.gradle.service.syncAction.impl.bridge.GradleBridgeModuleDataService"/>
+    <externalProjectDataService implementation="org.jetbrains.plugins.gradle.service.syncAction.impl.bridge.GradleBridgeFinalizerDataService"/>
 
     <postStartupActivity implementation="org.jetbrains.plugins.gradle.service.project.GradleStartupActivity"/>
     <backgroundPostStartupActivity implementation="org.jetbrains.plugins.gradle.service.project.GradleVersionUpdateStartupActivity"/>

plugins/gradle/plugin-resources/META-INF/plugin.xml

@@ -48,6 +48,8 @@ sealed interface GradleSyncPhase : Comparable<GradleSyncPhase> {
     }
   }
 
+  sealed interface DataServices: GradleSyncPhase
+
   companion object {
 
     /**
@@ -106,6 +108,9 @@ sealed interface GradleSyncPhase : Comparable<GradleSyncPhase> {
      */
     @JvmField
     val ADDITIONAL_MODEL_PHASE: GradleSyncPhase = GradleModelFetchPhase.ADDITIONAL_MODEL_PHASE.asSyncPhase()
+
+    @JvmField
+    val DATA_SERVICES_PHASE: GradleSyncPhase = GradleDataServicesSyncPhase()
   }
 }
 
@@ -122,7 +127,8 @@ private class GradleStaticSyncPhase(
     return when (other) {
       is GradleStaticSyncPhase -> order.compareTo(other.order)
       is GradleBaseScriptSyncPhase -> -1
-      is GradleDynamicSyncPhase -> -1
+      is GradleDynamicSyncPhase,
+      is GradleDataServicesSyncPhase -> -1
     }
   }
 
@@ -153,6 +159,7 @@ private class GradleDynamicSyncPhase(
       is GradleStaticSyncPhase -> 1
       is GradleBaseScriptSyncPhase -> 1
       is GradleDynamicSyncPhase -> modelFetchPhase.compareTo(other.modelFetchPhase)
+      is GradleDataServicesSyncPhase -> -1
     }
   }
 
@@ -178,7 +185,16 @@ private data object GradleBaseScriptSyncPhase: GradleSyncPhase.BaseScript {
     return when (other) {
       is GradleStaticSyncPhase -> 1
       is GradleBaseScriptSyncPhase -> 0
-      is GradleDynamicSyncPhase -> -1
+      is GradleDynamicSyncPhase,
+      is GradleDataServicesSyncPhase -> -1
     }
   }
+}
+
+private class GradleDataServicesSyncPhase: GradleSyncPhase.DataServices {
+
+  override val name: String = "DATA_SERVICES"
+
+  override fun compareTo(other: GradleSyncPhase): Int =
+    if (other is GradleDataServicesSyncPhase) 0 else 1
 }

plugins/gradle/src/org/jetbrains/plugins/gradle/service/syncAction/GradleSyncPhase.kt

@@ -0,0 +1,11 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package org.jetbrains.plugins.gradle.service.syncAction.impl.bridge
+
+import com.intellij.openapi.externalSystem.model.Key
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+object GradleBridgeFinalizerData {
+
+  val KEY: Key<GradleBridgeFinalizerData> = Key.create(GradleBridgeFinalizerData::class.java, Int.MAX_VALUE)
+}

plugins/gradle/src/org/jetbrains/plugins/gradle/service/syncAction/impl/bridge/GradleBridgeFinalizerData.kt

@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2025 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.jetbrains.plugins.gradle.service.syncAction.impl.bridge
+
+import com.intellij.openapi.externalSystem.model.DataNode
+import com.intellij.openapi.externalSystem.model.project.ProjectData
+import com.intellij.openapi.externalSystem.service.project.IdeModifiableModelsProvider
+import com.intellij.openapi.externalSystem.service.project.manage.AbstractProjectDataService
+import com.intellij.openapi.project.Project
+import com.intellij.openapi.util.registry.Registry
+import com.intellij.platform.workspace.storage.EntitySource
+import com.intellij.platform.workspace.storage.WorkspaceEntity
+import com.intellij.platform.workspace.storage.WorkspaceEntityBuilder
+import com.intellij.platform.workspace.storage.WorkspaceEntityWithSymbolicId
+import com.intellij.platform.workspace.storage.impl.WorkspaceEntityBase
+import org.jetbrains.plugins.gradle.service.syncAction.GradleEntitySource
+import org.jetbrains.plugins.gradle.service.syncAction.GradleSyncPhase
+
+
+class GradleBridgeFinalizerDataService : AbstractProjectDataService<GradleBridgeFinalizerData, Unit>() {
+  override fun getTargetDataKey() = GradleBridgeFinalizerData.KEY
+
+  override fun postProcess(toImport: Collection<DataNode<GradleBridgeFinalizerData?>?>,
+                           projectData: ProjectData?,
+                           project: Project,
+                           modelsProvider: IdeModifiableModelsProvider) {
+    if (!Registry.`is`("gradle.phased.sync.bridge.disabled") || projectData == null) return
+
+    val currentStorage = modelsProvider.actualStorageBuilder
+
+    val storageBeforeDataServices = modelsProvider.getUserData(SYNC_STORAGE_SNAPSHOT_BEFORE_DATA_SERVICES)!!
+    val index = storageBeforeDataServices.entitiesBySource {
+      sourceFilter(it, projectData.linkedExternalProjectPath)
+    }.associateBy {
+      WorkspaceEntityForLookup(it)
+    }
+
+    // Go over all the relevant entities and mark the ones that are not originally in the storage before data services execution
+    // with an explicit data service source. This is required because some entities otherwise inherit from their parents which are
+    // marked entity sources with explicit phases.
+    currentStorage.entitiesBySource {
+      sourceFilter(it, projectData.linkedExternalProjectPath)
+    }.filter {
+      if (it is WorkspaceEntityWithSymbolicId) {
+        storageBeforeDataServices.resolve(it.symbolicId)
+      } else {
+        index[WorkspaceEntityForLookup(it)]
+      } == null
+    }.forEach {
+      currentStorage.modifyEntity(WorkspaceEntityBuilder::class.java, it) {
+        entitySource = DataServiceEntitySource(projectData.linkedExternalProjectPath)
+      }
+    }
+  }
+
+  private fun sourceFilter(source: EntitySource, linkedExternalProjectPath: String) =
+    source is GradleEntitySource && source.projectPath == linkedExternalProjectPath
+
+  /** This is used for looking up entities without a symbolic ID. */
+  private class WorkspaceEntityForLookup(entity: WorkspaceEntity) {
+    val data = (entity as WorkspaceEntityBase).getData()
+
+    override fun equals(other: Any?): Boolean {
+      if (this === other) return true
+      if (javaClass != other?.javaClass) return false
+      return data.equalsByKey((other as WorkspaceEntityForLookup).data)
+    }
+
+    override fun hashCode() = data.hashCodeByKey()
+  }
+
+  private data class DataServiceEntitySource(override val projectPath: String): GradleEntitySource {
+    override val phase = GradleSyncPhase.DATA_SERVICES_PHASE
+  }
+}
+
+

plugins/gradle/src/org/jetbrains/plugins/gradle/service/syncAction/impl/bridge/GradleBridgeFinalizerDataService.kt

@@ -9,8 +9,13 @@ import com.intellij.openapi.externalSystem.service.project.manage.AbstractProjec
 import com.intellij.openapi.project.Project
 import com.intellij.openapi.util.registry.Registry
 import com.intellij.platform.workspace.jps.entities.ModuleEntity
+import com.intellij.platform.workspace.storage.ImmutableEntityStorage
 import com.intellij.platform.workspace.storage.entities
 import org.jetbrains.annotations.ApiStatus
+import com.intellij.openapi.util.Key as UserDataKey
+
+internal val SYNC_STORAGE_SNAPSHOT_BEFORE_DATA_SERVICES =
+  UserDataKey.create<ImmutableEntityStorage>("SYNC_STORAGE_SNAPSHOT_BEFORE_DATA_SERVICES")
 
 @ApiStatus.Internal
 class GradleBridgeProjectDataService : AbstractProjectDataService<GradleBridgeData, Unit>() {
@@ -26,6 +31,10 @@ class GradleBridgeProjectDataService : AbstractProjectDataService<GradleBridgeDa
     if (!Registry.`is`("gradle.phased.sync.bridge.disabled")) {
       removeModulesFromModelProvider(modelsProvider)
       removeEntitiesFromWorkspaceModel(modelsProvider)
+    } else {
+      // If the entities are not cleaned up, store a snapshot of the storage to be used for later post processing in.
+      // See [Int]
+      modelsProvider.putUserData(SYNC_STORAGE_SNAPSHOT_BEFORE_DATA_SERVICES, modelsProvider.actualStorageBuilder.toSnapshot())
     }
   }
 

plugins/gradle/src/org/jetbrains/plugins/gradle/service/syncAction/impl/bridge/GradleBridgeProjectDataService.kt

@@ -5,7 +5,9 @@ import com.intellij.gradle.toolingExtension.modelAction.GradleModelFetchPhase
 import com.intellij.openapi.observable.operation.OperationExecutionStatus
 import com.intellij.openapi.progress.ProgressManager
 import com.intellij.openapi.util.Disposer
+import com.intellij.openapi.util.registry.Registry
 import com.intellij.openapi.util.use
+import com.intellij.platform.backend.workspace.workspaceModel
 import com.intellij.platform.testFramework.assertion.WorkspaceAssertions
 import com.intellij.platform.testFramework.assertion.listenerAssertion.ListenerAssertion
 import com.intellij.platform.workspace.storage.toBuilder
@@ -14,6 +16,7 @@ import org.jetbrains.plugins.gradle.importing.TestModelProvider
 import org.jetbrains.plugins.gradle.importing.TestPhasedModel
 import org.jetbrains.plugins.gradle.service.project.DefaultProjectResolverContext
 import org.jetbrains.plugins.gradle.service.project.ProjectResolverContext
+import org.jetbrains.plugins.gradle.service.syncAction.GradleEntitySource
 import org.jetbrains.plugins.gradle.service.syncAction.GradleSyncPhase
 import org.jetbrains.plugins.gradle.service.syncAction.GradleSyncPhase.Dynamic.Companion.asSyncPhase
 import org.jetbrains.plugins.gradle.util.entity.GradleTestBridgeEntitySource
@@ -208,6 +211,7 @@ class GradlePhasedSyncTest : GradlePhasedSyncTestCase() {
                   true -> allDynamicPhases
                   else -> completedDynamicPhases
                 }
+                is GradleSyncPhase.DataServices -> error("Should not execute")
               }
               WorkspaceAssertions.assertEntities(myProject, expectedEntities.map { GradleTestEntityId(it) }) {
                 "Entities should be created for completed phases.\n" +
@@ -271,6 +275,7 @@ class GradlePhasedSyncTest : GradlePhasedSyncTestCase() {
               is GradleSyncPhase.Static -> completedStaticPhases
               is GradleSyncPhase.BaseScript -> completedBaseScriptPhases
               is GradleSyncPhase.Dynamic -> completedDynamicPhases
+              is GradleSyncPhase.DataServices -> error("Should not execute")
             }
             WorkspaceAssertions.assertEntities(myProject, expectedEntities.map { GradleTestEntityId(it) }) {
               "Bridge entities should be created for completed phases.\n" +
@@ -298,9 +303,101 @@ class GradlePhasedSyncTest : GradlePhasedSyncTestCase() {
         "Requested phases = $allPhases"
         "Completed phases = $completedPhases"
       }
+      val dataServicesEntities =  myProject.workspaceModel.currentSnapshot.entitiesBySource {
+        it is GradleEntitySource && it.phase == GradleSyncPhase.DATA_SERVICES_PHASE
+      }
+      Assertions.assertTrue(dataServicesEntities.toList().isEmpty()) {
+        "There should be no entities with phase ${GradleSyncPhase.DATA_SERVICES_PHASE}"
+      }
     }
   }
 
+  @Test
+  fun `test bridge entity contribution on Gradle sync phase with bridge disabled`() {
+    repeat(2) { index ->
+      Disposer.newDisposable().use { disposable ->
+        Registry.get("gradle.phased.sync.bridge.disabled").setValue(true, disposable)
+        val isSecondarySync = index == 1
+
+        val syncContributorAssertions = ListenerAssertion()
+        val syncPhaseCompletionAssertions = ListenerAssertion()
+
+        val allPhases = DEFAULT_SYNC_PHASES
+        val allStaticPhases = allPhases.filterIsInstance<GradleSyncPhase.Static>()
+        val allDynamicPhases = allPhases.filterIsInstance<GradleSyncPhase.Dynamic>()
+        val completedPhases = CopyOnWriteArrayList<GradleSyncPhase>()
+
+        for (phase in allPhases) {
+          addSyncContributor(phase, disposable) { context, storage ->
+            val builder = storage.toBuilder()
+            syncContributorAssertions.trace {
+              val entitySource = GradleTestEntitySource(context.projectPath, phase)
+              builder addEntity GradleTestEntity(phase, entitySource)
+              Assertions.assertTrue(completedPhases.add(phase)) {
+                "The $phase should be completed only once."
+              }
+            }
+            return@addSyncContributor builder.toSnapshot()
+          }
+          whenSyncPhaseCompleted(phase, disposable) { _ ->
+            syncPhaseCompletionAssertions.trace {
+              val completedStaticPhases = completedPhases.filterIsInstance<GradleSyncPhase.Static>()
+              val completedBaseScriptPhases = completedPhases.filterIsInstance<GradleSyncPhase.BaseScript>()
+              val completedDynamicPhases = completedPhases.filterIsInstance<GradleSyncPhase.Dynamic>()
+              val expectedEntities = when (phase) {
+                is GradleSyncPhase.Static -> when (isSecondarySync) {
+                  true -> completedStaticPhases + allDynamicPhases
+                  else -> completedStaticPhases
+                }
+                is GradleSyncPhase.BaseScript -> when (isSecondarySync) {
+                  true -> allStaticPhases + completedBaseScriptPhases + allDynamicPhases
+                  else -> completedBaseScriptPhases
+                }
+                is GradleSyncPhase.Dynamic -> when (isSecondarySync) {
+                  true -> allDynamicPhases
+                  else -> completedDynamicPhases
+                }
+                is GradleSyncPhase.DataServices -> error("Should not execute")
+              }
+              WorkspaceAssertions.assertEntities(myProject, expectedEntities.map { GradleTestEntityId(it) }) {
+                "Entities should be created for completed phases.\n" +
+                "Completed phases = $completedPhases\n"
+                "isSecondarySync = $isSecondarySync"
+              }
+            }
+          }
+        }
+
+        initMultiModuleProject()
+        importProject()
+        assertMultiModuleProjectStructure()
+
+        syncContributorAssertions.assertListenerFailures()
+        syncContributorAssertions.assertListenerState(allPhases.size) {
+          "All requested sync phases should be handled."
+        }
+        syncPhaseCompletionAssertions.assertListenerFailures()
+        syncPhaseCompletionAssertions.assertListenerState(allPhases.size) {
+          "All requested sync phases should be completed."
+        }
+
+        WorkspaceAssertions.assertEntities(myProject, allDynamicPhases.map { GradleTestEntityId(it) }) {
+          "Entities should be created for completed phases.\n" +
+          "Requested phases = $allPhases"
+          "Completed phases = $completedPhases"
+        }
+
+        val dataServicesEntities =  myProject.workspaceModel.currentSnapshot.entitiesBySource {
+          it is GradleEntitySource && it.phase == GradleSyncPhase.DATA_SERVICES_PHASE
+        }
+        Assertions.assertTrue(dataServicesEntities.toList().isNotEmpty()) {
+          "There should be at least one entity with phase ${GradleSyncPhase.DATA_SERVICES_PHASE}"
+        }
+      }
+    }
+  }
+
+
   @Test
   fun `test phased Gradle sync for custom static phase without model provider`() {
     `test phased Gradle sync for custom phase without model provider`(