[Cosmos] Clean up test databases left behind by CI runs

sdk/cosmos/AGENTS.md

@@ -0,0 +1,69 @@
+# AGENTS.md — Azure Cosmos DB SDK for Java
+
+This file provides Cosmos-specific guidance for AI agents working in the `sdk/cosmos` subtree.
+For general repo-wide guidance, see the [root AGENTS.md](../../AGENTS.md).
+
+## Test resource hygiene (required for any test that touches an account)
+
+Every live test stage in `tests.yml` now runs against a **long lived shared account** that is never
+torn down - either a fixed self-owned account in RG `sdk-ci` (the main and Http2 stages) or a thin
+client / GSI account. Several matrix legs and concurrent pipeline runs share them, so a database a
+test forgets to delete stays on that account forever.
+
+When adding or changing tests in `azure-cosmos-tests`:
+
+- **Create databases with `TestSuiteBase.createTestDatabase(client, "label")`**, never with
+  `UUID.randomUUID()`, a fixed literal, or a raw `client.createDatabase(...)` call. The helper embeds
+  the CI run id in the id so cleanup can attribute the database to this run without deleting another
+  run's in-flight resources. `TestResourceHygieneTest` fails the build on direct database creation.
+- **Still delete what you create** in an `@AfterClass(alwaysRun = true)` using `safeDeleteDatabase`.
+  `CosmosTestResourceJanitor` cleans up leftovers but **fails the run** when it has to.
+- Containers inside a database you delete need no separate cleanup.
+
+### How cleanup works
+
+| Layer | Where | Catches |
+| --- | --- | --- |
+| `CosmosTestResourceRegistry` + `CosmosTestResourceJanitor` | in the test JVM, at the end of the run | normal completion; deletes leftovers and **fails the run** naming the offending test |
+| JVM shutdown hook | in the test JVM | crashes and hard aborts |
+| `cleanup-test-resources.yml` post step | pipeline, `condition: always()` | failed jobs, and jobs whose JVM died. Runs on cancellation too, but only within `cancelTimeoutInMinutes` (5 by default), so treat it as best effort there |
+
+Every layer scopes deletion by the run id (`CosmosTestRunId`), derived from `System.JobId` (a GUID that
+is unique per job), with the build id carried along so a stray database can be traced back to a build.
+A run only ever deletes its own resources, so concurrent legs on a shared account cannot interfere.
+
+When a leak fails a run, TestNG has already written its reports by the time the janitor runs, so the
+**published test results show `Tests run: 0` rather than the leak** — which reads like an infrastructure
+fault. The leak is reported in the build log and as an ADO issue annotation; look there, not in the Tests
+tab.
+
+`CosmosTestAccountJanitor` is the standalone entry point used by the pipeline post step:
+
+```bash
+mvn -f sdk/cosmos/azure-cosmos-tests/pom.xml exec:java \
+  -Dexec.args="--account-host <uri> --account-key <key>"
+```
+
+It deletes only the current job's databases; pass `--run-id` to target another run.
+
+### Guardrails
+
+Two layers, because a static scan alone is not enough:
+
+- **At runtime**, `CosmosTestResourceRegistry` rejects any database id that CI cleanup could not
+  attribute to this run, failing the test immediately and naming it. This is the layer that actually
+  holds: it catches ids built at runtime, ids swapped inside a file that already has a ratchet
+  allowance, and creation through APIs the scanner does not know about. Disable only with
+  `-DCOSMOS.TEST_RESOURCE_ID_VALIDATION_ENABLED=false`.
+- **Statically**, `TestResourceHygieneTest` ratchets against direct database creation.
+  `azure-cosmos-tests/src/test/resources/test-resource-hygiene-baseline.properties` records the
+  violations that existed when the check was introduced, and the build fails when a file exceeds its
+  entry or a new offending file appears. It runs in the `unit` group, so it gives fast feedback on
+  every PR without needing an account. Do not add entries for new tests; when you migrate a file,
+  lower or remove its entry.
+
+### Escape hatches
+
+- `-DCOSMOS.TEST_RESOURCE_JANITOR_ENABLED=false` disables the in-process janitor entirely.
+- `-DCOSMOS.TEST_RESOURCE_JANITOR_FAIL_ON_LEAK=false` keeps the cleanup but stops it from failing the
+  run. Use this only to unblock a pipeline while the leaking test is being fixed.

sdk/cosmos/azure-cosmos-spark-account-data-resolver-sample/pom.xml

@@ -375,6 +375,9 @@
               <source>1.8</source>
               <target>1.8</target>
               <scalaVersion>2.12.19</scalaVersion>
+              <!-- Per-module bridge cache. The default (~/.sbt/1.0/zinc/org.scala-sbt) is shared, and the
+                   plugin installs the bridge with no lock, so parallel builds (-T) can read a half-written jar. -->
+              <secondaryCacheDir>${project.build.directory}/scala-compiler-bridge</secondaryCacheDir>
             </configuration>
             <executions>
               <execution>

sdk/cosmos/azure-cosmos-spark_3/pom.xml

@@ -429,6 +429,9 @@
               <source>${maven.compiler.source}</source>
               <target>${maven.compiler.target}</target>
               <scalaVersion>${scala.version}</scalaVersion>
+              <!-- Per-module bridge cache. The default (~/.sbt/1.0/zinc/org.scala-sbt) is shared, and the
+                   plugin installs the bridge with no lock, so parallel builds (-T) can read a half-written jar. -->
+              <secondaryCacheDir>${project.build.directory}/scala-compiler-bridge</secondaryCacheDir>
             </configuration>
             <executions>
               <execution>

sdk/cosmos/azure-cosmos-tests/pom.xml

@@ -256,6 +256,18 @@ Licensed under the MIT License.
         </configuration>
       </plugin>
 
+      <!-- Runs CosmosTestAccountJanitor, the standalone cleanup entry point used by the always-run
+           pipeline post steps. -->
+      <plugin>
+        <groupId>org.codehaus.mojo</groupId>
+        <artifactId>exec-maven-plugin</artifactId>
+        <version>3.5.1</version> <!-- {x-version-update;org.codehaus.mojo:exec-maven-plugin;external_dependency} -->
+        <configuration>
+          <mainClass>com.azure.cosmos.CosmosTestAccountJanitor</mainClass>
+          <classpathScope>test</classpathScope>
+        </configuration>
+      </plugin>
+
       <!-- CosmosSkip - Needed temporary false values to not fail. -->
       <plugin>
         <groupId>org.apache.maven.plugins</groupId>

sdk/cosmos/azure-cosmos-tests/src/test/java/com/azure/cosmos/CosmosDatabaseForTest.java

@@ -3,22 +3,222 @@
 
 package com.azure.cosmos;
 
+import com.azure.cosmos.models.CosmosDatabaseProperties;
+import com.azure.cosmos.models.SqlParameter;
+import com.azure.cosmos.models.SqlQuerySpec;
+import com.azure.cosmos.util.CosmosPagedFlux;
+import org.apache.commons.lang3.RandomStringUtils;
+import org.apache.commons.lang3.StringUtils;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.time.LocalDateTime;
 import java.time.ZoneOffset;
 import java.time.format.DateTimeFormatter;
-import java.util.UUID;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.function.Predicate;
 
+/**
+ * Naming and lifetime helper for databases created by the test suite.
+ * <p>
+ * Ids follow {@code RxJava.SDKTest.SharedDatabase_<timestamp>_<runId>_<random>}. The run id lets cleanup
+ * delete exactly the databases created by the current run without touching resources that a concurrently
+ * executing matrix leg or pipeline run is still using on the same shared account. The legacy three
+ * segment form (without a run id) is still parsed, so those ids are recognized as test databases and
+ * never mistaken for hand created fixtures.
+ * <p>
+ * Timestamps are UTC: an id is written by the test agent and may be compared on a different machine.
+ */
 public final class CosmosDatabaseForTest {
+    private static final Logger LOGGER = LoggerFactory.getLogger(CosmosDatabaseForTest.class);
     public static final String SHARED_DB_ID_PREFIX = "RxJava.SDKTest.SharedDatabase";
     private static final String DELIMITER = "_";
     private static final DateTimeFormatter TIME_FORMATTER = DateTimeFormatter.ofPattern("yyyyMMdd'T'HHmmss");
+    private static final int RANDOM_SUFFIX_LENGTH = 8;
+    private static final int MAX_LABEL_LENGTH = 10;
 
     private CosmosDatabaseForTest() {
     }
 
+    private static LocalDateTime nowUtc() {
+        return LocalDateTime.now(ZoneOffset.UTC);
+    }
+
     public static String generateId() {
-        return SHARED_DB_ID_PREFIX + DELIMITER + TIME_FORMATTER.format(LocalDateTime.now(ZoneOffset.UTC))
-            + DELIMITER + UUID.randomUUID();
+        return generateId(null);
+    }
+
+    /**
+     * Generates a run tagged database id. The optional label only makes logs and portal views readable -
+     * uniqueness comes from the timestamp plus the random suffix, and cleanup scoping from the run id.
+     * <p>
+     * Length matters and is deliberately kept short. The database id is embedded in derived resource
+     * names elsewhere in the SDK - notably the throughput control group id, which is
+     * {@code <database>/<container>/<group>/<suffix>}, base64 encoded and then extended with a UUID to
+     * form a control item id. A long database id overflows that and the emulator rejects the request
+     * with "400 Bad Request - Invalid URL", so keep this comfortably shorter than the ~82 character ids
+     * that are known to work.
+     *
+     * @param label optional human readable label, may be null.
+     * @return a database id that cleanup is able to attribute to the current run.
+     */
+    public static String generateId(String label) {
+        String random = RandomStringUtils.randomAlphanumeric(RANDOM_SUFFIX_LENGTH);
+        String suffix = StringUtils.isEmpty(label) ? random : sanitizeLabel(label) + random;
+
+        return SHARED_DB_ID_PREFIX
+            + DELIMITER + TIME_FORMATTER.format(nowUtc())
+            + DELIMITER + CosmosTestRunId.get()
+            + DELIMITER + suffix;
+    }
+
+    private static String sanitizeLabel(String label) {
+        String sanitized = label.replaceAll("[^A-Za-z0-9]", "");
+        return sanitized.length() <= MAX_LABEL_LENGTH ? sanitized : sanitized.substring(0, MAX_LABEL_LENGTH);
+    }
+
+    /**
+     * @param id the database id to check.
+     * @return true when the id was produced by {@link #generateId(String)} - and is therefore owned by
+     * the test suite and safe for automated cleanup to delete.
+     */
+    public static boolean isTestDatabaseId(String id) {
+        return parse(id) != null;
+    }
+
+    private static ParsedId parse(String id) {
+        if (id == null) {
+            return null;
+        }
+
+        String[] parts = StringUtils.split(id, DELIMITER);
+        // 3 parts: legacy <prefix>_<timestamp>_<random>. 4 parts: <prefix>_<timestamp>_<runId>_<random>.
+        if (parts == null || parts.length < 3 || parts.length > 4) {
+            return null;
+        }
+        if (!StringUtils.equals(parts[0], SHARED_DB_ID_PREFIX)) {
+            return null;
+        }
+
+        try {
+            // Parsed purely to validate the id: a malformed timestamp means this is not one of ours, so
+            // it is left alone. The value itself is not needed - deletion is only ever run scoped.
+            LocalDateTime.parse(parts[1], TIME_FORMATTER);
+            String runId = parts.length == 4 ? parts[2] : null;
+            return new ParsedId(runId);
+        } catch (Exception e) {
+            return null;
+        }
+    }
+
+    /**
+     * Deletes every database created by the currently executing run, regardless of age. This is the
+     * safety net for tests that created a database and failed to delete it.
+     *
+     * @param client the database manager to clean up with.
+     * @return what the sweep deleted, and whether it ran to completion.
+     */
+    public static CleanupResult cleanupDatabasesForCurrentRun(DatabaseManager client) {
+        return cleanupDatabasesForRun(client, CosmosTestRunId.get());
+    }
+
+    /**
+     * Deletes every database created by the given run, regardless of age. Used by the pipeline post step,
+     * which knows the run id of the job that just finished.
+     *
+     * @param client the database manager to clean up with.
+     * @param runId the run id to delete databases for.
+     * @return what the sweep deleted, and whether it ran to completion.
+     */
+    public static CleanupResult cleanupDatabasesForRun(DatabaseManager client, String runId) {
+        if (StringUtils.isEmpty(runId)) {
+            // Matching a null run id would match every legacy (pre run id) database, which may belong to a
+            // run that is still executing.
+            throw new IllegalArgumentException("runId must not be empty");
+        }
+
+        LOGGER.info("Cleaning test databases for run {} ...", runId);
+        return deleteMatching(client, dbForTest -> StringUtils.equals(dbForTest.runId, runId));
+    }
+
+    private static CleanupResult deleteMatching(DatabaseManager client, Predicate<ParsedId> predicate) {
+        List<String> deleted = new ArrayList<>();
+        int failures = 0;
+
+        List<CosmosDatabaseProperties> dbs = client.queryDatabases(
+            new SqlQuerySpec(
+                "SELECT * FROM c WHERE STARTSWITH(c.id, @PREFIX)",
+                Collections.singletonList(new SqlParameter("@PREFIX", SHARED_DB_ID_PREFIX))))
+            .collectList()
+            .block();
+
+        if (dbs == null) {
+            return new CleanupResult(deleted, failures);
+        }
+
+        for (CosmosDatabaseProperties db : dbs) {
+            ParsedId parsed = parse(db.getId());
+            // A null parsed id means the id does not follow the test convention - it may be a hand created
+            // fixture, so leave it alone.
+            if (parsed == null || !predicate.test(parsed)) {
+                continue;
+            }
+
+            LOGGER.info("Deleting database {}", db.getId());
+            try {
+                client.getDatabase(db.getId()).delete().block();
+                deleted.add(db.getId());
+            } catch (Exception e) {
+                // Keep going - one undeletable database must not strand the rest - but remember that this
+                // sweep did not fully succeed, so callers do not mistake "found nothing" for "all clean".
+                failures++;
+                LOGGER.warn("Failed to delete database {}", db.getId(), e);
+            } finally {
+                CosmosTestResourceRegistry.unregisterDatabase(db.getId());
+            }
+        }
+
+        return new CleanupResult(deleted, failures);
+    }
+
+    /**
+     * What a cleanup sweep managed to do. A sweep that deleted nothing because every delete failed is not
+     * the same as a sweep that found nothing, and callers must be able to tell them apart.
+     */
+    public static final class CleanupResult {
+        private final List<String> deletedDatabaseIds;
+        private final int failureCount;
+
+        private CleanupResult(List<String> deletedDatabaseIds, int failureCount) {
+            this.deletedDatabaseIds = deletedDatabaseIds;
+            this.failureCount = failureCount;
+        }
+
+        public List<String> getDeletedDatabaseIds() {
+            return this.deletedDatabaseIds;
+        }
+
+        public int getFailureCount() {
+            return this.failureCount;
+        }
+
+        public boolean isComplete() {
+            return this.failureCount == 0;
+        }
     }
 
+    private static final class ParsedId {
+        private final String runId;
+
+        private ParsedId(String runId) {
+            this.runId = runId;
+        }
+    }
+
+    public interface DatabaseManager {
+        CosmosPagedFlux<CosmosDatabaseProperties> queryDatabases(SqlQuerySpec query);
+        CosmosAsyncDatabase getDatabase(String id);
+    }
 }