Docker group browse: read members' images by name, not in OFFSET pages

crates/db/src/oci_dao.rs

@@ -1329,6 +1329,75 @@ pub async fn list_images(
     .await
 }
 
+/// The next `limit` images of a repository whose name contains `query` (all
+/// when empty), by name, after the image `after` (from the first when
+/// empty): a keyset scan over the `(repository_id, image, digest)` index,
+/// then each image's counts from its own manifests (a lookup per image, so
+/// that no plan reads the repository's other manifests). Reading a
+/// repository's images batch by batch costs each batch alone, not the
+/// batches before it as `OFFSET` does.
+pub async fn list_images_after<'e, E>(
+    executor: E,
+    repository_id: i64,
+    query: &str,
+    after: &str,
+    limit: i64,
+) -> sqlx::Result<Vec<ImageSummary>>
+where
+    E: sqlx::PgExecutor<'e>,
+{
+    let query = query.to_ascii_lowercase();
+    sqlx::query_as!(
+        ImageSummary,
+        r#"
+        WITH names AS (
+            SELECT DISTINCT image FROM oci_manifest
+            WHERE repository_id = $1 AND image > $3 AND strpos(image, $2) > 0
+            ORDER BY image
+            LIMIT $4
+        )
+        SELECT n.image AS "image!",
+               s.tag_count AS "tag_count!",
+               s.manifest_count AS "manifest_count!",
+               s.last_pushed_at AS "last_pushed_at!"
+        FROM names n
+        CROSS JOIN LATERAL (
+            SELECT count(DISTINCT t.tag) AS tag_count,
+                   count(DISTINCT m.id) AS manifest_count,
+                   max(m.created_at) AS last_pushed_at
+            FROM oci_manifest m
+            LEFT JOIN oci_tag t ON t.manifest_id = m.id
+            WHERE m.repository_id = $1 AND m.image = n.image
+        ) s
+        ORDER BY n.image
+        "#,
+        repository_id,
+        query,
+        after,
+        limit,
+    )
+    .fetch_all(executor)
+    .await
+}
+
+/// Which of `images` one of the repositories `repository_ids` holds.
+pub async fn images_held(
+    pool: &PgPool,
+    repository_ids: &[i64],
+    images: &[String],
+) -> sqlx::Result<Vec<String>> {
+    sqlx::query_scalar!(
+        r#"
+        SELECT DISTINCT image AS "image!" FROM oci_manifest
+        WHERE repository_id = ANY($1) AND image = ANY($2)
+        "#,
+        repository_ids,
+        images,
+    )
+    .fetch_all(pool)
+    .await
+}
+
 /// How many images [`list_images`] would list in total.
 pub async fn count_images(pool: &PgPool, repository_id: i64, query: &str) -> sqlx::Result<i64> {
     let query = query.to_ascii_lowercase();

crates/db/tests/oci_images.rs

@@ -0,0 +1,98 @@
+//! `oci_dao`'s image listings against a real Postgres: reading every image
+//! of a repository a batch at a time by name (#596).
+
+mod common;
+
+use forge_db::oci_dao;
+use forge_db::repository_dao::{self, NewRepository};
+
+use common::{pool, unique};
+
+/// Index entries and table rows of `oci_manifest` this transaction read.
+async fn rows_read(conn: &mut sqlx::PgConnection) -> i64 {
+    sqlx::query_scalar(
+        r#"
+        SELECT coalesce(sum(pg_stat_get_xact_tuples_returned(i.indexrelid)), 0)::bigint
+             + coalesce(pg_stat_get_xact_tuples_returned('oci_manifest'::regclass), 0)
+        FROM pg_index i
+        WHERE i.indrelid = 'oci_manifest'::regclass
+        "#,
+    )
+    .fetch_one(conn)
+    .await
+    .unwrap()
+}
+
+/// Reading a repository's images batch by batch with `list_images_after`
+/// reads each manifest a bounded number of times, however many batches came
+/// before: an `OFFSET` page reads every image before it again, so reading
+/// them all costs the square of their number.
+#[tokio::test]
+async fn reading_every_image_by_name_reads_each_manifest_a_few_times() {
+    let pool = pool().await;
+    let repo = repository_dao::insert_repository(
+        &pool,
+        NewRepository::hosted(format!("images-{}", unique())),
+    )
+    .await
+    .unwrap()
+    .id;
+    let sha = format!("{:064x}", (u128::from(unique()) << 64) | 0x596);
+    sqlx::query(
+        "INSERT INTO blob (sha256, size, s3_key, state) VALUES ($1, 1, $1, 'committed') \
+         ON CONFLICT DO NOTHING",
+    )
+    .bind(&sha)
+    .execute(&pool)
+    .await
+    .unwrap();
+    const IMAGES: i64 = 10_000;
+    sqlx::query(
+        r#"
+        INSERT INTO oci_manifest (repository_id, image, digest, blob_sha256, media_type, size)
+        SELECT $1, 'team/app-' || lpad(n::text, 5, '0'), 'sha256:' || $2, $2,
+               'application/vnd.oci.image.manifest.v1+json', 1
+        FROM generate_series(0, $3 - 1) AS n
+        "#,
+    )
+    .bind(repo)
+    .bind(&sha)
+    .bind(IMAGES)
+    .execute(&pool)
+    .await
+    .unwrap();
+    sqlx::query("ANALYZE oci_manifest")
+        .execute(&pool)
+        .await
+        .unwrap();
+
+    let mut tx = pool.begin().await.unwrap();
+    let before = rows_read(&mut tx).await;
+    let mut listed = Vec::new();
+    let mut after = String::new();
+    loop {
+        let batch = oci_dao::list_images_after(&mut *tx, repo, "app", &after, 1000)
+            .await
+            .unwrap();
+        let full = batch.len() == 1000;
+        let Some(last) = batch.last() else {
+            break;
+        };
+        after = last.image.clone();
+        listed.extend(batch.into_iter().map(|image| image.image));
+        if !full {
+            break;
+        }
+    }
+    let read = rows_read(&mut tx).await - before;
+    assert_eq!(listed.len() as i64, IMAGES);
+    assert!(
+        listed.windows(2).all(|w| w[0] < w[1]),
+        "in order, once each"
+    );
+    assert_eq!(listed[0], "team/app-00000");
+    assert!(
+        read < 3 * IMAGES,
+        "listing {IMAGES} images read {read} manifest rows"
+    );
+}

crates/server/src/browse_images.rs

@@ -329,6 +329,12 @@ async fn image_page(
 /// group's content selector applies, and its guard decides per image which
 /// members serve it (the registry's `GroupMembers`). An image is listed
 /// with the counts of the highest priority member serving it.
+///
+/// Each leaf's images are read a batch at a time by name (a keyset scan,
+/// [`oci_dao::list_images_after`]), so a listing costs what the members
+/// hold, once. Whether an owner holds an image is asked only of the images
+/// a leaf that is no owner would list, a batch at a time: an owner's other
+/// images are never read.
 async fn group_image_page(
     state: &AppState,
     principal: &Principal,
@@ -345,55 +351,87 @@ async fn group_image_page(
         .group_members(principal, group, 0)
         .await
         .map_err(failure_response)?;
-    // The images an owner has, whatever the caller may read: with
-    // `prefer_hosted`, only owners serve them.
-    let mut owned: HashSet<String> = HashSet::new();
-    if members.asks_owners() {
-        for owner in members.owner_ids() {
-            for_each_image(state, owner, &query.q, |image| {
-                owned.insert(image.image);
-            })
-            .await
-            .map_err(db_error)?;
-        }
-    }
+    let owners = members.owner_ids();
     let mut merged: BTreeMap<String, oci_dao::ImageSummary> = BTreeMap::new();
     for (position, leaf) in members.leaves() {
-        for_each_image(state, leaf.id, &query.q, |image| {
-            if merged.contains_key(&image.image) || !group_filter.allows(&image.image) {
-                return;
+        // An owner serves what it holds whoever else holds it.
+        let asks_owners = members.asks_owners() && !owners.contains(&leaf.id);
+        let mut after = String::new();
+        loop {
+            let batch = oci_dao::list_images_after(
+                &state.pool,
+                leaf.id,
+                &query.q,
+                &after,
+                IMAGE_BATCH_SIZE,
+            )
+            .await
+            .map_err(db_error)?;
+            let full = batch.len() as i64 == IMAGE_BATCH_SIZE;
+            let Some(last) = batch.last() else {
+                break;
+            };
+            after = last.image.clone();
+            let candidates: Vec<oci_dao::ImageSummary> = batch
+                .into_iter()
+                .filter(|image| {
+                    !merged.contains_key(&image.image) && group_filter.allows(&image.image)
+                })
+                .collect();
+            // Owned or not makes a difference only where nothing else
+            // (a reserved prefix) keeps the image to the owners already.
+            let ask: Vec<String> = candidates
+                .iter()
+                .filter(|image| asks_owners && !members.owners_only(&image.image, false))
+                .map(|image| image.image.clone())
+                .collect();
+            let owned: HashSet<String> = if ask.is_empty() {
+                HashSet::new()
+            } else {
+                oci_dao::images_held(&state.pool, &owners, &ask)
+                    .await
+                    .map_err(db_error)?
+                    .into_iter()
+                    .collect()
+            };
+            for image in candidates {
+                let owners_only = members.owners_only(&image.image, owned.contains(&image.image));
+                if members.serves(position, &image.image, owners_only) {
+                    merged.insert(image.image.clone(), image);
+                }
             }
-            let owners_only = members.owners_only(&image.image, owned.contains(&image.image));
-            if members.serves(position, &image.image, owners_only) {
-                merged.insert(image.image.clone(), image);
+            if !full {
+                break;
             }
-        })
-        .await
-        .map_err(db_error)?;
+        }
     }
     let total = merged.len();
     let images = merged.into_values().skip(query.skip).take(take).collect();
     Ok((total, images))
 }
 
 /// Call `each` with every image of repository `repository_id` matching
-/// `q`, read [`IMAGE_BATCH_SIZE`] at a time.
+/// `q`, read [`IMAGE_BATCH_SIZE`] at a time by name.
 async fn for_each_image(
     state: &AppState,
     repository_id: i64,
     q: &str,
     mut each: impl FnMut(oci_dao::ImageSummary),
 ) -> sqlx::Result<()> {
-    let mut offset = 0;
+    let mut after = String::new();
     loop {
         let images =
-            oci_dao::list_images(&state.pool, repository_id, q, offset, IMAGE_BATCH_SIZE).await?;
-        let count = images.len() as i64;
+            oci_dao::list_images_after(&state.pool, repository_id, q, &after, IMAGE_BATCH_SIZE)
+                .await?;
+        let full = images.len() as i64 == IMAGE_BATCH_SIZE;
+        let Some(last) = images.last() else {
+            return Ok(());
+        };
+        after = last.image.clone();
         images.into_iter().for_each(&mut each);
-        if count < IMAGE_BATCH_SIZE {
+        if !full {
             return Ok(());
         }
-        offset += count;
     }
 }