From d93749c7d79931c27f46fbc2ae21837d4b7a16e6 Mon Sep 17 00:00:00 2001 From: JasonWildMe Date: Sun, 28 Jun 2026 12:30:48 -0700 Subject: [PATCH] feat(token API): kNN /similarity endpoint + concurrency guard + request-scoped PITs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses the critical scaling bottleneck from live API feedback (downloading every 2152-float vector to find matches) AND the operator requirement that high-volume agents must not interfere with interactive Wildbook use. POST /api/v3/similarity (token, read-only): server-side k-NN re-identification — given an annotation the caller can see, return the most similar annotations with scores, WITHOUT shipping any vectors. Surfaces Wildbook's existing OpenSearch kNN. - Filtered ANN: ACL + same viewpoint + method/version + optional taxonomy/locationId scope + self/encounter exclusion all go INSIDE the native knn.filter (an outer bool.filter would post-filter ANN results). For non-admin, also an outer ACL wrap (defense in depth). Filtered nested kNN verified empirically on the deployed cluster. - Never returns vectors (_source excludes embeddings + defensive remove()). - No-oracle visibility gate on the query annotation (not-visible == not-found == 404). - Deterministic single-embedding resolution; viewpoint required (400 if null); k 1-100. - Re-ID metric wildbook_token_reid_queries_total, distinct from pipeline ID gauges. Concurrency guard (TokenApiConcurrency): per-instance semaphores bound how many token requests run at once so volume is unlimited but live users aren't starved. Wired into ALL token endpoints: SIMILARITY pool (default 2) for the heavy kNN; GENERAL pool (default 8) for search + media/resolve. AutoCloseable permit (no leak/over-release); best-effort metrics; instant 429 + Retry-After; inflight + rejected Prometheus gauges. Sizes via OpenSearch.properties (restart to resize). Request-scoped PITs (OpenSearch.queryPitPrivate): the similarity endpoint never touches the shared static PIT_CACHE, so a token call can't delete/race a PIT an interactive UI search is using. (Migrating existing token search/media-resolve to private PITs is the recommended fast-follow.) Specs Codex design-reviewed before implementation; code Codex-reviewed. Tests (46 across the touched suites): TokenApiConcurrencyTest (3); SimilarityApiTest (9: auth/ validation, embeddings-stripped, ambiguous-embedding 400, filtered-ANN shape + self-exclusion, no-oracle 404, 429 guard); MediaResolveApi + SearchApi GENERAL-guard 429 tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../specs/2026-06-28-similarity-endpoint.md | 160 ++++++++++ .../2026-06-28-token-api-concurrency-guard.md | 131 +++++++++ src/main/java/org/ecocean/OpenSearch.java | 61 +++- .../java/org/ecocean/api/MediaResolveApi.java | 14 + src/main/java/org/ecocean/api/SearchApi.java | 18 +- .../java/org/ecocean/api/SimilarityApi.java | 276 ++++++++++++++++++ .../org/ecocean/api/TokenApiConcurrency.java | 111 +++++++ src/main/webapp/WEB-INF/web.xml | 11 + .../org/ecocean/api/MediaResolveApiTest.java | 31 ++ .../org/ecocean/api/SimilarityApiTest.java | 270 +++++++++++++++++ .../ecocean/api/TokenApiConcurrencyTest.java | 67 +++++ 11 files changed, 1144 insertions(+), 6 deletions(-) create mode 100644 docs/superpowers/specs/2026-06-28-similarity-endpoint.md create mode 100644 docs/superpowers/specs/2026-06-28-token-api-concurrency-guard.md create mode 100644 src/main/java/org/ecocean/api/SimilarityApi.java create mode 100644 src/main/java/org/ecocean/api/TokenApiConcurrency.java create mode 100644 src/test/java/org/ecocean/api/SimilarityApiTest.java create mode 100644 src/test/java/org/ecocean/api/TokenApiConcurrencyTest.java diff --git a/docs/superpowers/specs/2026-06-28-similarity-endpoint.md b/docs/superpowers/specs/2026-06-28-similarity-endpoint.md new file mode 100644 index 0000000000..8922430d28 --- /dev/null +++ b/docs/superpowers/specs/2026-06-28-similarity-endpoint.md @@ -0,0 +1,160 @@ +# Token Similarity (k-NN) Endpoint — Design + +**Date:** 2026-06-28 +**Branch:** `feature/similarity-endpoint` +**Driver:** Live API feedback (MMF/Sharkbook) — the **critical** scaling bottleneck. To find matches an +agent must download every catalog annotation's 2,152-float vector (≈0.5–0.9 s/annotation; a 200-record +page ≈176 s; >10k unwalkable). Wildbook already does OpenSearch kNN internally for matching +(`Annotation.getMatchQuery`, nested `knn` on `embeddings.vector` filtered by method/version). Surface +that as a read-only token endpoint so the agent gets nearest neighbors **without downloading vectors**. + +## Endpoint +`POST /api/v3/similarity` — token-auth (same `tokenAuthSearch` Shiro filter as search), read-only. A new +path (not under `/api/v3/search/*`, which the SearchApi servlet owns). + +### Request +```json +{ "annotationId": "", "k": 50, "methodVersion": "msv4.1", + "taxonomy": "Equus quagga", "locationId": "Mbuluzi Game Reserve", "excludeSameEncounter": true } +``` +- `annotationId` (**required**) — the query annotation. The server reads ITS stored embedding; the + caller never sends a vector (v1 has no raw-vector input — avoids untrusted vectors and is the whole + point: no vector download). +- `k` (optional, default 50, **max 100**) — number of neighbors. +- `method` (optional, default `miewid`) / `methodVersion` (optional) — which embedding to compare. If + omitted and the annotation has exactly one (method,version) embedding, use it; if it has several and + `methodVersion` is omitted → 400 (ambiguous). +- `taxonomy` (optional) / `locationId` (optional) — explicit scope filters, applied as exact term + filters on `encounterTaxonomy` / `encounterLocationId` inside the candidate filter. **v1 uses these + explicit, safe scope params rather than an arbitrary caller `query`** — this avoids needing to + validate a free-form ranking-oracle filter (deferred to a future version with a strict allow-list). +- `excludeSameEncounter` (optional, default false) — also drop candidates sharing the query's + `encounterId` (the skills' anti-inflation rule). + +### Behavior +1. Resolve the query annotation. If the caller can't see it (ACL) or it doesn't exist → **404**, the + two indistinguishable (no existence oracle, mirroring media/resolve). +2. Get its embedding for (method, methodVersion). None → **400** `"no embedding for that + method/version"`. +3. Build the nested `knn` query (vector + `k`) filtered by `embeddings.method`/`methodVersion` + (reuse the `getMatchQuery` construction). +4. **Enforce correctness server-side:** candidates are constrained to the **same `viewpoint`** as the + query annotation (skip this constraint only if the query's viewpoint is null), and the same + method/version (already in the knn filter). Always **exclude the query annotation itself**; if + `excludeSameEncounter`, also exclude its `encounterId`. +5. **ACL-scope candidates:** wrap the candidate query with `applyAclFilter` for non-admin (admins + unfiltered), so neighbors are only annotations the caller may see. +6. Merge the optional caller `query` as an additional `filter`. +7. Execute via `queryPit` with size = k; return the top-k. + +### Response +```json +{ "query": { "annotationId": "...", "viewpoint": "left", "method": "miewid", "methodVersion": "msv4.1" }, + "neighbors": [ { "id": "...", "score": 0.87, "encounterId": "...", "individualId": "...", + "viewpoint": "left", "iaClass": "zebra", "methodVersion": ["msv4.1"] }, ... ] } +``` +- Each neighbor is the **sanitized** annotation doc (ACL-scrubbed, like search hits) **with the + `embeddings` field stripped** + a `score` (the OpenSearch kNN cosine similarity). Stripping + `embeddings` is essential: it keeps the payload tiny (the whole point) and never returns vectors. +- `score` lets the agent threshold/rank (calibrate per the skill guidance). + +## Security +- Token-auth, read-only; candidates ACL-scoped via `applyAclFilter`; query annotation must be visible + (else 404, no oracle). +- **Never returns vectors** (embeddings stripped from neighbor `_source`). +- `k` bounded (≤100); optional `query` validated (no scripts/aggs/runtime_mappings); same-viewpoint + + same-method/version enforced (prevents cross-latent-space comparison the skills warn against). +- Self (and optionally same-encounter) excluded so the top hit isn't the query itself. + +## Metrics (REQUIRED — this is user-run re-ID compute on our servers) +A similarity call is a re-identification query executed on our OpenSearch. It creates **no `ia.Task`**, +so the existing identification gauges (`MetricsBot` JDOQL counts of `Task` rows with +`ibeis.identification`/`pipeline_root`/`graph`) will NOT see it. Add a dedicated metric so this re-ID +channel is visible in `/metrics` and impact reporting: +- A Prometheus **Counter** on `io.prometheus.client.CollectorRegistry.defaultRegistry` (the same + registry `MetricsResource`/`WildbookMetrics` expose), e.g. `wildbook_token_reid_queries_total` + (label: `context`), incremented once per successfully-served similarity query. Optionally also a + neighbors-returned counter and a latency histogram. +- Keep it **distinct** from the pipeline `wildbook_identification_tasks*` gauges — token-API re-ID is a + separate, synchronous, agent-driven modality; conflating the two corrupts identification counts (cf. + the v2-detection-gauge pollution incident). Increment only on a real query that ran the kNN (not on + 400/404/validation rejections). +- Counter resets on restart (standard Prometheus semantics; the scrape/TSDB accumulates). If a durable + all-time total is also wanted, additionally persist a running total (e.g. a `SystemValue`) — optional. + +## Reuse / files +- New servlet `org.ecocean.api.SimilarityApi extends ApiBase` (`doPost`); web.xml servlet + url-pattern + `/api/v3/similarity`; Shiro `[urls]` `/api/v3/similarity = tokenAuthSearch`. +- New `OpenSearch.knnQuery(...)` (or reuse `Annotation.getMatchQuery` building blocks) to assemble the + nested-knn + filters; reuse `applyAclFilter`, `queryPit`, `sanitizeDoc`. +- Strip `embeddings` from each returned `_source` before `sanitizeDoc`/output. + +## Out of scope (future) +- Raw-vector input; cross-viewpoint search; paging beyond top-k; multi-annotation batch; returning the + query vector. A `/api/v3/agent-skill/find-missed-matches` rewrite to use this endpoint (huge speedup) + is a follow-up doc PR once this lands. + +## Design review (BINDING — incorporated) + +**PREREQUISITE — RESOLVED (verified empirically on zebra.wildme.org 2026-06-28).** Through the existing +`/api/v3/search/annotation` (which forwards the body to OpenSearch), a nested `knn` with an inner +`filter` is accepted (HTTP 200) and the filter is **honored**: filtering to the opposite `viewpoint` +returned 10 `left` neighbors instead of the natural `right` ones, and the filter worked on BOTH a nested +subfield (`embeddings.methodVersion`) and a PARENT field (`viewpoint`). So the deployed engine supports +efficient filtered nested kNN; build on native `knn.filter` (parent + nested filters), no mapping +migration needed. **Heavy-compute note:** kNN is the most expensive token operation — it must be behind +the token-API concurrency guard (see the separate request-isolation workstream) so it cannot starve +interactive Wildbook search/matching. + +(Original prerequisite text, now satisfied:) verify the deployed engine supports efficient *filtered +nested* kNN. The +`embeddings.vector` mapping is `knn_vector` with **no `method.engine`** specified, so filtered-kNN +support is engine/version-dependent and unconfirmed. The endpoint MUST NOT ship on a global-ANN + +outer-`bool.filter` (post-filter) basis — that returns fewer than `k`, lets hidden/out-of-scope vectors +suppress visible neighbors, and is a recall/leak hazard. Resolve one of: + (a) confirm the live OpenSearch honors `knn.filter` on the nested `embeddings.vector` (empirically: + POST a `knn`+`filter` body through the existing `/api/v3/search/annotation` and check it filters); + (b) migrate the mapping to a filter-capable engine (Lucene) + reindex; + (c) fall back to exact filtered scoring or a flattened embedding-row index. +Do not implement until (a) or (b) holds. + +**Corrected query shape:** put ACL, viewpoint, method/version, caller scope, and self/same-encounter +exclusion **inside the native `knn.filter`** (the efficient filtered-ANN path), NOT only an outer bool. +Keep an outer ACL `bool.filter` as defense-in-depth only. + +**Vector non-leak (defense in depth):** set `_source.excludes:["embeddings"]` on the request, reject any +caller `_source`/`fields`/`docvalue_fields`, AND `remove("embeddings")` from each doc before output. + +**Query-annotation visibility gating (no-oracle):** mirror `MediaResolveApi` exactly — token context +only; parse first; for a non-admin run the singleton id-eligibility ACL query (fresh PIT) to confirm +visibility BEFORE any DB/vector load; not-visible and not-found return an identical **404** (never a +distinct "no embedding"/timing signal for hidden ids). + +**Caller `query` is now a ranking oracle → strict validation** (stricter than annotation search): allow +filter-only DSL on a **field allow-list**; reject scripts, aggs, runtime_mappings, `_source`, `fields`, +sort, post_filter, rescore, `nested`/`embeddings`/`vector`, ACL fields, terms-lookup, and overly +large/deep clauses (extend the existing `DENY_FEATURES`/`nodeAllowed` deny-list). + +**Deterministic method resolution:** do NOT reuse `getEmbeddingByMethod(null,null)` (returns first of a +set). Require exactly one eligible `(method, methodVersion)` after defaults (method=miewid); else 400. + +**Viewpoint-null → 400** (don't silently broaden latent space). **k:** default 50, range 1–100, else 400. + +**Cost:** fresh PIT; reject huge/deep caller filters; note token/IP rate-limiting as a platform follow-up. + +## Open questions (RESOLVED by review) +- Path `/api/v3/similarity` (not under `/api/v3/search/*`). • 404 for not-visible/not-found query + annotation. • Always the query annotation's viewpoint (no override) in v1; null → 400. • k default 50, + max 100, out-of-range → 400. • Do NOT rely on outer `bool.filter`; use native `knn.filter`. + +## Original open questions (superseded by the resolutions above) +1. Endpoint path/name: `/api/v3/similarity` vs `/api/v3/search/similar` (the latter collides with the + SearchApi `/api/v3/search/*` mapping → avoid). Confirm `/api/v3/similarity`. +2. Not-visible/not-found query annotation: 404 (chosen) vs 200-empty — which better preserves no-oracle + while being ergonomic? +3. Should same-viewpoint enforcement be overridable via an explicit `viewpoint` param, or always the + query's viewpoint? (Lean: always the query's in v1; overridable later.) +4. `k` max (100 chosen) and default (50) — reasonable for re-ID shortlists? +5. Does the OpenSearch nested-knn query honor an outer `bool.filter` (ACL + viewpoint + scope) so the + ANN search is restricted to the filtered candidate set (efficient + correct), or does kNN run global + then post-filter? (Implementation must confirm the filter applies to the ANN candidate set.) diff --git a/docs/superpowers/specs/2026-06-28-token-api-concurrency-guard.md b/docs/superpowers/specs/2026-06-28-token-api-concurrency-guard.md new file mode 100644 index 0000000000..9e9752bacc --- /dev/null +++ b/docs/superpowers/specs/2026-06-28-token-api-concurrency-guard.md @@ -0,0 +1,131 @@ +# Token-API Concurrency Guard — Design + +**Date:** 2026-06-28 +**Branch:** `feature/similarity-endpoint` (ships under PR5, ahead of/with the kNN endpoint) +**Driver:** Operator requirement — a remote agent may make a *massive volume* of token-API calls (that's +fine), but it **must not interfere with interactive Wildbook use**. Token search/kNN hit the SAME +OpenSearch the live UI and the live re-ID matching pipeline use; one sequential agent already triggered +intermittent `500`s under sustained paging. kNN (PR5) is the heaviest op. We need **resource isolation**, +not a volume quota. + +## Approach: bound *concurrency*, not volume +A per-instance (per-JVM) **semaphore** caps how many token-API requests execute *simultaneously*. The +agent may send unlimited requests; only N run at once; the rest get **`429 Too Many Requests` + +`Retry-After`** and back off. This bounds the agent's instantaneous footprint on OpenSearch/Tomcat, +leaving headroom for live users, without limiting total work (the agent self-paces). Volume is fine; +crowding-out is prevented. + +### Two pools +- **GENERAL** — every *token* search / `media/resolve` / aggregation request acquires one permit. +- **SIMILARITY** — the kNN `/api/v3/similarity` endpoint acquires from a **separate, smaller** pool + (kNN is CPU-heaviest). Independent pools (one acquire per request) — no double-acquire, no deadlock. +Defaults (per instance, tunable): `GENERAL = 16`, `SIMILARITY = 4`, `Retry-After = 5s`. Read via the +existing `getConfigurationValue(key, default)` mechanism (e.g. `tokenApiMaxConcurrent`, +`tokenApiMaxConcurrentSimilarity`, `tokenApiRetryAfterSeconds`); fall back to defaults. + +### Only token requests are throttled +The guard applies **only when the request is token-authenticated** (`TOKEN_AUTH_ATTR`). Interactive +session/UI traffic (the live users we're protecting) is never throttled. Non-token requests don't +consume a permit. + +### Mechanism: in-servlet, not a filter +Each token-serving servlet acquires/releases around its work (robust, testable, no Shiro-ordering +dependence): +``` +if (tokenAuth) { + if (!TokenApiConcurrency.tryAcquire(kind)) { write 429 + Retry-After; return; } + try { ...handle... } finally { TokenApiConcurrency.release(kind); } +} +``` +- `tryAcquire` is **non-blocking** (immediate) — never hold a Tomcat thread waiting (waiting threads are + themselves a way to starve the pool). Acquire as late as possible but before the expensive + OpenSearch/DB work (`SearchApi`: before `queryPit`; `MediaResolveApi`: before the resolve loop; + `SimilarityApi`: at entry, before the kNN). Release in `finally` on every path (success, 4xx, 5xx, + exception) so permits never leak. +- Wiring: `SearchApi` + `MediaResolveApi` → `Kind.GENERAL`; `SimilarityApi` → `Kind.SIMILARITY`. + +### 429 response +`HTTP 429`, header `Retry-After: `, body `{"error":"token API busy, retry after N seconds"}`. +No `_source`/data. Distinct from `403`/`401`/`400`. + +### Metrics (ties into the re-ID metrics work) +On the Prometheus `defaultRegistry`: +- `wildbook_token_api_inflight{pool=}` — gauge of permits currently held (inc on acquire, dec on + release). +- `wildbook_token_api_rejected_total{pool=}` — counter of `429`s. +This makes saturation observable (ops can see when agents are hitting the cap and tune the budget). + +## Security / correctness +- Permits released in `finally` (no leak on exceptions/early returns). +- Per-JVM semaphore: in a fleet, each instance self-protects (load-balancer spreads load; per-instance + caps still bound each node's contribution to the shared OpenSearch). A distributed limiter is overkill. +- Fail-open vs fail-closed on guard error: the guard itself must never 500 a request; if anything in the + guard throws, proceed (fail-open) — the guard is a safety valve, not an auth gate. +- Does not change auth/ACL; orthogonal to the per-request cost ceilings (10k window, k≤100, etc.) which + remain and bound each permit's cost. + +## Out of scope (documented) +- Distributed/cluster-wide limiting; a dedicated OpenSearch read replica / lower-priority search pool + (the stronger infra isolation) — recommended as a platform follow-up, not in this PR. +- Per-token / per-IP fairness (one noisy token vs many) — v1 is a global per-instance cap; per-token + fairness can come later if needed. +- Short-wait acquire (vs instant 429) — start instant; revisit if 429 churn is high. + +## Testing +- `TokenApiConcurrency`: acquire up to N succeeds; the N+1th fails; release restores a permit; pools are + independent (exhausting SIMILARITY doesn't block GENERAL). Config override changes sizes; bad config → + defaults. Metrics inflight inc/dec and rejected increments. +- Servlet-level: with the pool exhausted (acquire all permits in the test), a token request returns + `429` + `Retry-After` and does NOT call `queryPit`/resolve; a session (non-token) request is NOT + throttled even when the pool is exhausted; permit released after a normal request (next succeeds); + permit released after an exception path. + +## Design review — incorporated (BINDING) + +- **Scope honesty:** this is *bounded concurrency* for OpenSearch isolation, not absolute isolation. + It does not protect Tomcat threads / Postgres / body-parsing from a hostile flood (rejected token + requests still parse + do user lookup before the 429). Document a WAF/edge rate-limit + an OpenSearch + `thread_pool`/search-backpressure layer as the complementary platform follow-ups. +- **AutoCloseable permit handle (no leak / no over-release):** `tryAcquire(kind)` returns a `Permit` + (`AutoCloseable`, or null if not acquired). Use `try (Permit p = …) { … }`; never call bare + `Semaphore.release()` (over-releases if acquire failed). Handles `MediaResolveApi`'s IOException + rethrow safely. +- **Acquire placement (cover ALL OpenSearch work, skip 4xx gates):** + - `SearchApi`: acquire AFTER the individual-token 400 gate, BEFORE `new OpenSearch()`/`deletePit`/ + `queryPit` (so 401/403/404/405/400/agg-400 never consume a permit, but the whole OpenSearch path is + covered). + - `MediaResolveApi`: acquire AFTER token/context/body/`currentUser`/`isAdmin`, BEFORE computing + `visible` (which itself calls `gatedVisibleIds` → `deletePit`+`queryPit`). + - `SimilarityApi`: at entry to the kNN execution. +- **Lower defaults** (one sequential agent already caused 500s): `GENERAL = 8`, `SIMILARITY = 2` + (configurable). Note the budgets are **additive** (≤10 concurrent token OpenSearch ops/instance). + Clamp invalid config to ≥1; **restart required** to resize; these are node/webapp-wide + `OpenSearch.properties` settings read once via `getConfigurationValue`. +- **Instant 429** (no wait); GENERAL pool also covers aggregation requests in v1 (their per-request cost + is already capped: terms-only, size≤1000). +- **Idempotent metrics registration:** register the gauge/counter once (guard against duplicate + registration across redeploys/multiple webapp contexts — e.g. catch the already-registered case or use + a singleton holder). Update the inflight gauge only through the permit lifecycle. Low-cardinality + labels (`pool` only). + +## Related interference vector — shared PIT cache (IMPORTANT, partly addressed here) +`OpenSearch.PIT_CACHE` is a **static `HashMap` keyed by index**, shared by token AND interactive/UI +search. Each token request calls `deletePit(indexName)`, which deletes the cached point-in-time the UI +may be mid-search on — so a single token call can disrupt a live search, independent of concurrency. The +unsynchronized HashMap is also a cross-thread hazard. +- **In this PR:** the new `SimilarityApi` uses a **request-scoped PIT** (create → use → delete its own + id; never touches `PIT_CACHE`), so the heaviest new endpoint adds no shared-PIT contention. +- **Follow-up PR (recommended, possibly urgent — affects the already-built token search/`media/resolve`):** + move those token paths to request-scoped PITs too (and/or make `PIT_CACHE` concurrency-safe). Done as a + focused PR because it touches shared search infra used by the live UI. The concurrency guard only + *bounds* how many token requests churn the shared PIT at once; it does not stop a single one from + deleting a UI PIT. + +## Open questions (for review) +1. Defaults: `GENERAL=16`, `SIMILARITY=4` per instance — reasonable starting headroom, or derive from + CPU count? (Lean: fixed configurable defaults; ops tune per instance size.) +2. Instant `429` vs a short bounded wait (e.g. 100ms) before rejecting — which better avoids 429 churn + without risking thread pile-up? (Lean: instant.) +3. Should aggregation requests share the GENERAL pool or get their own (they can be heavier than a plain + term search)? (Lean: GENERAL for v1; split if needed.) +4. Per-JVM semaphore acceptable given the fleet, or is per-token/per-IP fairness needed in v1? diff --git a/src/main/java/org/ecocean/OpenSearch.java b/src/main/java/org/ecocean/OpenSearch.java index 1ed22fb2ad..3a1aeb7c78 100644 --- a/src/main/java/org/ecocean/OpenSearch.java +++ b/src/main/java/org/ecocean/OpenSearch.java @@ -413,6 +413,48 @@ public JSONObject queryPit(String indexName, final JSONObject query, int numFrom return new JSONObject(rtn); } + /** + * Request-scoped PIT search: create a PRIVATE point-in-time, run the query, and always delete that + * PIT. Unlike {@link #queryPit}, it never reads or writes the shared static {@code PIT_CACHE}, so a + * one-shot token API call cannot delete or race a PIT an interactive/UI search is using. Use this + * on the token path (one query per call); does not mutate the caller's {@code query}. + */ + public JSONObject queryPitPrivate(String indexName, final JSONObject query, int numFrom, + int pageSize, String sort, String sortOrder) + throws IOException { + if (!isValidIndexName(indexName)) throw new IOException("invalid index name: " + indexName); + Request createReq = new Request("POST", + indexName + "/_search/point_in_time?keep_alive=" + SEARCH_PIT_TIME); + String pitId = new JSONObject(getRestResponse(createReq)).optString("pit_id", null); + if (pitId == null) throw new IOException("failed to get PIT id"); + try { + JSONObject q = new JSONObject(query.toString()); // copy: do not mutate caller's query + q.put("from", numFrom); + q.put("size", pageSize); + if (sort != null) { + JSONArray sortArr = new JSONArray(); + if ((sortOrder == null) || !"desc".equals(sortOrder)) sortOrder = "asc"; + sortArr.put(new JSONObject("{\"" + sort + "\":{\"order\":\"" + sortOrder + "\"}}")); + q.put("sort", sortArr); + } + JSONObject jpit = new JSONObject(); + jpit.put("id", pitId); + jpit.put("keep_alive", SEARCH_PIT_TIME); + q.put("pit", jpit); + Request searchRequest = new Request("POST", "/_search?track_total_hits=true"); + searchRequest.setJsonEntity(q.toString()); + return new JSONObject(getRestResponse(searchRequest)); + } finally { + try { // best-effort delete of the private PIT + Request del = new Request("DELETE", "/_search/point_in_time"); + del.setJsonEntity(new JSONObject().put("pit_id", pitId).toString()); + getRestResponse(del); + } catch (Exception ignore) { + System.out.println("queryPitPrivate(): best-effort PIT delete failed: " + ignore); + } + } + } + // just return the actual hit results // note: each object in the array has _id but actual doc is in _source!! public static JSONArray getHits(JSONObject queryResults) { @@ -963,10 +1005,12 @@ public static JSONObject applyEncounterAclFilter(JSONObject query, String userId return applyAclFilter(query, userId, "encounter"); } - public static JSONObject applyAclFilter(JSONObject query, String userId, String indexName) - throws IOException { - if ((query == null) || !Util.stringExists(userId)) - throw new IOException("applyAclFilter: null query or userId"); + /** + * The ACL access clause as a standalone bool (publiclyReadable OR submitter OR viewUser), for + * embedding directly inside a query — e.g. inside a native {@code knn.filter} so the kNN candidate + * set itself is ACL-scoped (an OUTER bool.filter would post-filter the ANN results instead). + */ + public static JSONObject aclShouldClause(String userId, String indexName) { // encounter docs carry a single submitterUserId; annotation/individual carry the union set submitterUserIds String submitterField = "encounter".equals(indexName) ? "submitterUserId" : "submitterUserIds"; JSONArray should = new JSONArray(); @@ -976,7 +1020,14 @@ public static JSONObject applyAclFilter(JSONObject query, String userId, String JSONObject aclBool = new JSONObject(); aclBool.put("should", should); aclBool.put("minimum_should_match", 1); - JSONObject acl = new JSONObject().put("bool", aclBool); + return new JSONObject().put("bool", aclBool); + } + + public static JSONObject applyAclFilter(JSONObject query, String userId, String indexName) + throws IOException { + if ((query == null) || !Util.stringExists(userId)) + throw new IOException("applyAclFilter: null query or userId"); + JSONObject acl = aclShouldClause(userId, indexName); JSONObject wrapBool = new JSONObject(); JSONObject inner = query.optJSONObject("query"); diff --git a/src/main/java/org/ecocean/api/MediaResolveApi.java b/src/main/java/org/ecocean/api/MediaResolveApi.java index 08ef9ed1fe..10d63c9b79 100644 --- a/src/main/java/org/ecocean/api/MediaResolveApi.java +++ b/src/main/java/org/ecocean/api/MediaResolveApi.java @@ -81,6 +81,7 @@ private void handle(HttpServletRequest request, HttpServletResponse response) } Shepherd myShepherd = null; + TokenApiConcurrency.Permit permit = null; try { myShepherd = new Shepherd(context); myShepherd.setAction("api.MediaResolveApi.POST"); @@ -91,6 +92,10 @@ private void handle(HttpServletRequest request, HttpServletResponse response) return; } boolean isAdmin = currentUser.isAdmin(myShepherd); + // Concurrency guard: bound simultaneous token OpenSearch work so a high-volume agent + // cannot starve interactive Wildbook use. Acquire before the gate/resolve OpenSearch calls. + permit = TokenApiConcurrency.tryAcquire(TokenApiConcurrency.Kind.GENERAL); + if (permit == null) { writeBusy(response); return; } Set visible = isAdmin ? ids : gatedVisibleIds(ids, currentUser.getId()); JSONArray results = new JSONArray(); for (String id : visible) { @@ -105,10 +110,19 @@ private void handle(HttpServletRequest request, HttpServletResponse response) writeError(response, 500, "resolve failed"); ex.printStackTrace(); } finally { + if (permit != null) permit.close(); if (myShepherd != null) myShepherd.rollbackAndClose(); } } + private void writeBusy(HttpServletResponse response) throws IOException { + response.setStatus(429); + response.setHeader("Retry-After", Integer.toString(TokenApiConcurrency.RETRY_AFTER_SECONDS)); + writeBody(response, new JSONObject().put("error", + "token API busy, retry after " + TokenApiConcurrency.RETRY_AFTER_SECONDS + " seconds") + .toString()); + } + private void writeError(HttpServletResponse response, int status, String msg) throws IOException { response.setStatus(status); writeBody(response, new JSONObject().put("error", msg).toString()); diff --git a/src/main/java/org/ecocean/api/SearchApi.java b/src/main/java/org/ecocean/api/SearchApi.java index 58f60e4fba..49f2d97227 100644 --- a/src/main/java/org/ecocean/api/SearchApi.java +++ b/src/main/java/org/ecocean/api/SearchApi.java @@ -156,8 +156,21 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) System.out.println("SearchApi search indexName=" + indexName + " tokenAuth=" + tokenAuth); - OpenSearch os = new OpenSearch(); + // Concurrency guard (token only): bound simultaneous token OpenSearch work so a + // high-volume agent cannot starve interactive Wildbook use. Session/UI unthrottled. + TokenApiConcurrency.Permit permit = tokenAuth + ? TokenApiConcurrency.tryAcquire(TokenApiConcurrency.Kind.GENERAL) : null; + if (tokenAuth && (permit == null)) { + response.setStatus(429); + response.setHeader("Retry-After", + Integer.toString(TokenApiConcurrency.RETRY_AFTER_SECONDS)); + res.put("success", false); + res.put("error", "token API busy, retry after " + + TokenApiConcurrency.RETRY_AFTER_SECONDS + " seconds"); + } else { try { + // construct INSIDE the try so a client-init failure still releases the permit + OpenSearch os = new OpenSearch(); if (deletePit) os.deletePit(indexName); JSONObject queryRes = os.queryPit(indexName, query, numFrom, pageSize, sort, sortOrder); @@ -194,7 +207,10 @@ protected void doPost(HttpServletRequest request, HttpServletResponse response) res.put("success", false); res.put("error", "query failed"); ex.printStackTrace(); + } finally { + if (permit != null) permit.close(); } + } // end concurrency-guard else } // end individual-token field-gate else } } diff --git a/src/main/java/org/ecocean/api/SimilarityApi.java b/src/main/java/org/ecocean/api/SimilarityApi.java new file mode 100644 index 0000000000..d383ca53d8 --- /dev/null +++ b/src/main/java/org/ecocean/api/SimilarityApi.java @@ -0,0 +1,276 @@ +package org.ecocean.api; + +import io.prometheus.client.Counter; +import java.io.IOException; +import java.util.LinkedHashSet; +import java.util.Set; +import javax.servlet.ServletException; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import org.ecocean.Annotation; +import org.ecocean.Embedding; +import org.ecocean.Encounter; +import org.ecocean.OpenSearch; +import org.ecocean.User; +import org.ecocean.Util; +import org.ecocean.security.WildbookTokenAuthenticationFilter; +import org.ecocean.servlet.ServletUtilities; +import org.ecocean.shepherd.core.Shepherd; +import org.json.JSONArray; +import org.json.JSONObject; + +/** + * POST /api/v3/similarity — read-only, token-authenticated k-NN re-identification: given an annotation + * the caller can see, return the most visually similar annotations (server-side ANN), WITHOUT the + * caller downloading any vectors. Surfaces Wildbook's existing OpenSearch kNN as a bounded, ACL-scoped + * token endpoint. + * + * Safety: candidates are filtered INSIDE the native knn.filter (ACL + same viewpoint + method/version + * + optional scope + self/encounter exclusion) so the ANN candidate set itself is constrained (an + * outer bool.filter would post-filter ANN results). Vectors are never returned. The query annotation's + * visibility is gated like media/resolve (not-visible and not-found are the same 404 — no existence + * oracle). It runs on a request-scoped PIT (no shared PIT_CACHE contention) and behind the SIMILARITY + * concurrency pool so it cannot starve interactive Wildbook use. Each served query increments a + * dedicated re-ID metric (this is user-run identification compute on our servers). + */ +public class SimilarityApi extends ApiBase { + + static final int DEFAULT_K = 50; + static final int MAX_K = 100; + + // Distinct from the pipeline wildbook_identification_tasks* gauges: token-API re-ID is a separate, + // synchronous, agent-driven modality. Registered defensively (re-deploy / multi-context safe). + private static final Counter REID_QUERIES; + static { + Counter c = null; + try { + c = Counter.build().name("wildbook_token_reid_queries_total") + .help("Token-API similarity (k-NN re-identification) queries served, by context") + .labelNames("context").register(); + } catch (RuntimeException alreadyRegistered) { + c = null; + } + REID_QUERIES = c; + } + + @Override + protected void doPost(HttpServletRequest request, HttpServletResponse response) + throws ServletException, IOException { + handle(request, response); + } + + void doPostForTest(HttpServletRequest request, HttpServletResponse response) throws IOException { + handle(request, response); + } + + private void handle(HttpServletRequest request, HttpServletResponse response) throws IOException { + boolean tokenAuth = Boolean.TRUE.equals( + request.getAttribute(WildbookTokenAuthenticationFilter.TOKEN_AUTH_ATTR)); + if (!tokenAuth) { writeError(response, 401, "token auth required"); return; } + String context = (String) request.getAttribute( + WildbookTokenAuthenticationFilter.TOKEN_CONTEXT_ATTR); + if (!Util.stringExists(context)) { writeError(response, 401, "unauthorized"); return; } + + // Parse + validate BEFORE acquiring a permit or touching the DB/OpenSearch. + String annotationId, method, methodVersion, taxonomy, locationId; + int k; + boolean excludeSameEncounter; + try { + JSONObject body = ServletUtilities.jsonFromHttpServletRequest(request); + if (body == null) { writeError(response, 400, "missing request body"); return; } + annotationId = body.optString("annotationId", null); + if (!Util.stringExists(annotationId)) { + writeError(response, 400, "annotationId is required"); return; + } + k = body.has("k") ? body.optInt("k", -1) : DEFAULT_K; + if ((k < 1) || (k > MAX_K)) { + writeError(response, 400, "k must be between 1 and " + MAX_K); return; + } + method = body.optString("method", "miewid"); + if (!Util.stringExists(method)) method = "miewid"; + methodVersion = Util.stringExists(body.optString("methodVersion", null)) + ? body.optString("methodVersion") : null; + taxonomy = Util.stringExists(body.optString("taxonomy", null)) + ? body.optString("taxonomy") : null; + locationId = Util.stringExists(body.optString("locationId", null)) + ? body.optString("locationId") : null; + excludeSameEncounter = body.optBoolean("excludeSameEncounter", false); + } catch (Exception ex) { + writeError(response, 400, "malformed request body"); return; + } + + // Concurrency guard: bound simultaneous kNN so token re-ID cannot starve interactive use. + TokenApiConcurrency.Permit permit = + TokenApiConcurrency.tryAcquire(TokenApiConcurrency.Kind.SIMILARITY); + if (permit == null) { writeBusy(response); return; } + + Shepherd myShepherd = null; + try { + myShepherd = new Shepherd(context); + myShepherd.setAction("api.SimilarityApi.POST"); + myShepherd.beginDBTransaction(); + User currentUser = myShepherd.getUser(request); + if ((currentUser == null) || (currentUser.getId() == null)) { + writeError(response, 401, "unauthorized"); return; + } + boolean isAdmin = currentUser.isAdmin(myShepherd); + + // No-oracle visibility gate for the QUERY annotation: a non-admin must be able to see it. + // not-visible and not-found both -> identical 404 (no existence oracle). + if (!isAdmin) { + Set one = new LinkedHashSet<>(); + one.add(annotationId); + Set visible = visibleAnnotationIds(one, currentUser.getId()); + if (!visible.contains(annotationId)) { writeError(response, 404, "not found"); return; } + } + Annotation ann = myShepherd.getAnnotation(annotationId); + if (ann == null) { writeError(response, 404, "not found"); return; } + + String viewpoint = ann.getViewpoint(); + if (!Util.stringExists(viewpoint)) { + writeError(response, 400, "query annotation has no viewpoint"); return; + } + Embedding emb = resolveSingleEmbedding(ann, method, methodVersion); + if (emb == null) { + writeError(response, 400, "no unique embedding for that method/version"); return; + } + String mv = emb.getMethodVersion(); + String queryEncounterId = null; + Encounter qEnc = ann.findEncounter(myShepherd); + if (qEnc != null) queryEncounterId = qEnc.getId(); + + JSONObject body = buildKnnBody(emb, k, method, mv, viewpoint, taxonomy, locationId, + annotationId, (excludeSameEncounter ? queryEncounterId : null), + (isAdmin ? null : currentUser.getId())); + // Defense in depth: the ACL is already inside knn.filter (the candidate set is ACL-scoped); + // for a non-admin also wrap the whole query so any hit is ACL-gated even if the engine ever + // post-filtered. (No-op for admin.) + if (!isAdmin) body = OpenSearch.applyAclFilter(body, currentUser.getId(), "annotation"); + + OpenSearch os = new OpenSearch(); + JSONObject queryRes = os.queryPitPrivate("annotation", body, 0, k, null, null); + JSONArray neighbors = new JSONArray(); + JSONObject outerHits = queryRes.optJSONObject("hits"); + JSONArray hits = (outerHits != null) ? outerHits.optJSONArray("hits") : null; + if (hits != null) { + for (int i = 0; i < hits.length(); i++) { + JSONObject h = hits.optJSONObject(i); + if (h == null) continue; + JSONObject doc = h.optJSONObject("_source"); + if (doc == null) continue; + doc.remove("embeddings"); // never return vectors (defense in depth on top of _source excludes) + JSONObject clean = OpenSearch.sanitizeDoc(doc, "annotation", myShepherd, currentUser, + tokenAuth); + clean.put("score", h.optDouble("_score", 0.0)); + neighbors.put(clean); + } + } + + JSONObject q = new JSONObject(); + q.put("annotationId", annotationId); + q.put("viewpoint", viewpoint); + q.put("method", method); + q.put("methodVersion", mv); + JSONObject out = new JSONObject(); + out.put("query", q); + out.put("neighbors", neighbors); + if (REID_QUERIES != null) REID_QUERIES.labels(context).inc(); + response.setStatus(200); + writeBody(response, out.toString()); + } catch (IOException ioEx) { + throw ioEx; // broken pipe etc — don't double-write + } catch (Exception ex) { + writeError(response, 500, "similarity failed"); + ex.printStackTrace(); + } finally { + permit.close(); + if (myShepherd != null) myShepherd.rollbackAndClose(); + } + } + + /** Exactly one embedding matching (method, methodVersion-if-given); else null (ambiguous/none). */ + private Embedding resolveSingleEmbedding(Annotation ann, String method, String methodVersion) { + if (ann.getEmbeddings() == null) return null; + Embedding found = null; + for (Embedding e : ann.getEmbeddings()) { + if (e == null) continue; + if ((method != null) && !method.equals(e.getMethod())) continue; + if ((methodVersion != null) && !methodVersion.equals(e.getMethodVersion())) continue; + if (found != null) return null; // more than one match -> ambiguous + found = e; + } + return found; + } + + /** Build the nested-kNN body with ALL constraints inside the native knn.filter (filtered ANN). */ + private JSONObject buildKnnBody(Embedding emb, int k, String method, String methodVersion, + String viewpoint, String taxonomy, String locationId, String selfId, String excludeEncounterId, + String aclUserId) { + JSONArray filt = new JSONArray(); + filt.put(term("embeddings.method", method)); + filt.put(term("embeddings.methodVersion", methodVersion)); + filt.put(term("viewpoint", viewpoint)); + if (taxonomy != null) filt.put(term("encounterTaxonomy", taxonomy)); + if (locationId != null) filt.put(term("encounterLocationId", locationId)); + if (aclUserId != null) filt.put(OpenSearch.aclShouldClause(aclUserId, "annotation")); + JSONArray mustNot = new JSONArray(); + mustNot.put(term("id", selfId)); // never return the query annotation itself + if (excludeEncounterId != null) mustNot.put(term("encounterId", excludeEncounterId)); + JSONObject filterBool = new JSONObject().put("filter", filt).put("must_not", mustNot); + + JSONObject knnInner = new JSONObject() + .put("vector", new JSONArray(emb.vectorToFloatArray())) + .put("k", k) + .put("filter", new JSONObject().put("bool", filterBool)); + JSONObject nested = new JSONObject().put("nested", new JSONObject() + .put("path", "embeddings") + .put("query", new JSONObject().put("knn", + new JSONObject().put("embeddings.vector", knnInner)))); + JSONObject body = new JSONObject().put("query", nested); + // keep payload small AND never ship vectors: exclude embeddings from _source + body.put("_source", new JSONObject().put("excludes", new JSONArray().put("embeddings"))); + return body; + } + + private static JSONObject term(String field, String value) { + return new JSONObject().put("term", new JSONObject().put(field, value)); + } + + /** Which of these annotation ids the user may see, via a request-scoped (private) PIT. */ + private Set visibleAnnotationIds(Set ids, String userId) throws IOException { + JSONObject query = OpenSearch.buildIdEligibilityQuery(ids); + query = OpenSearch.applyAclFilter(query, userId, "annotation"); + JSONObject res = new OpenSearch().queryPitPrivate("annotation", query, 0, ids.size(), null, null); + Set visible = new LinkedHashSet<>(); + JSONObject outer = (res != null) ? res.optJSONObject("hits") : null; + JSONArray hits = (outer != null) ? outer.optJSONArray("hits") : null; + if (hits != null) { + for (int i = 0; i < hits.length(); i++) { + JSONObject h = hits.optJSONObject(i); + String hid = (h != null) ? h.optString("_id", null) : null; + if (Util.stringExists(hid)) visible.add(hid); + } + } + return visible; + } + + private void writeBusy(HttpServletResponse response) throws IOException { + response.setStatus(429); + response.setHeader("Retry-After", Integer.toString(TokenApiConcurrency.RETRY_AFTER_SECONDS)); + writeBody(response, new JSONObject() + .put("error", "token API busy, retry after " + TokenApiConcurrency.RETRY_AFTER_SECONDS + + " seconds").toString()); + } + + private void writeError(HttpServletResponse response, int status, String msg) throws IOException { + response.setStatus(status); + writeBody(response, new JSONObject().put("error", msg).toString()); + } + + private void writeBody(HttpServletResponse response, String s) throws IOException { + response.setContentType("application/json; charset=UTF-8"); + java.io.PrintWriter w = response.getWriter(); + w.write(s); + w.close(); + } +} diff --git a/src/main/java/org/ecocean/api/TokenApiConcurrency.java b/src/main/java/org/ecocean/api/TokenApiConcurrency.java new file mode 100644 index 0000000000..8afc68f79c --- /dev/null +++ b/src/main/java/org/ecocean/api/TokenApiConcurrency.java @@ -0,0 +1,111 @@ +package org.ecocean.api; + +import io.prometheus.client.Counter; +import io.prometheus.client.Gauge; +import java.util.concurrent.Semaphore; +import org.ecocean.OpenSearch; + +/** + * Per-instance concurrency guard for the read-only token API. A remote agent may make an unlimited + * VOLUME of calls — but only a bounded number execute concurrently, so token traffic cannot starve + * interactive Wildbook users (token search/kNN share the same OpenSearch as the live UI and the live + * re-ID matching pipeline). Excess token requests get HTTP 429 + Retry-After and back off. + * + * Two independent pools: GENERAL (token search / media-resolve / aggregations) and a smaller + * SIMILARITY (kNN — the heaviest op). Only token-authenticated requests acquire a permit; interactive + * session/UI traffic (the users we are protecting) is never throttled. + * + * Sizes are node/webapp-wide settings (OpenSearch.properties), read once at class load — restart to + * resize: tokenApiMaxConcurrent (default 8), tokenApiMaxConcurrentSimilarity (default 2), + * tokenApiRetryAfterSeconds (default 5). Budgets are additive (up to general+similarity concurrent ops + * per instance). This bounds OpenSearch contention; it is not a Tomcat/DB flood guard (see the spec). + */ +public final class TokenApiConcurrency { + + public enum Kind { GENERAL, SIMILARITY } + + private static final int GENERAL_MAX = + atLeast1((Integer) OpenSearch.getConfigurationValue("tokenApiMaxConcurrent", 8)); + private static final int SIMILARITY_MAX = + atLeast1((Integer) OpenSearch.getConfigurationValue("tokenApiMaxConcurrentSimilarity", 2)); + public static final int RETRY_AFTER_SECONDS = + atLeast1((Integer) OpenSearch.getConfigurationValue("tokenApiRetryAfterSeconds", 5)); + + private static final Semaphore GENERAL_SEM = new Semaphore(GENERAL_MAX); + private static final Semaphore SIMILARITY_SEM = new Semaphore(SIMILARITY_MAX); + + // Saturation observability. Registered defensively: a re-registration (redeploy / multiple webapp + // contexts in one JVM) must not fail class init, so on failure we run without these metrics. + private static final Gauge INFLIGHT; + private static final Counter REJECTED; + static { + Gauge g = null; + Counter c = null; + try { + g = Gauge.build().name("wildbook_token_api_inflight") + .help("Token API requests currently executing, by pool").labelNames("pool").register(); + c = Counter.build().name("wildbook_token_api_rejected_total") + .help("Token API requests rejected with 429 because the pool was at capacity, by pool") + .labelNames("pool").register(); + } catch (RuntimeException alreadyRegistered) { + g = null; + c = null; + } + INFLIGHT = g; + REJECTED = c; + } + + private TokenApiConcurrency() {} + + private static int atLeast1(Integer v) { + return ((v == null) || (v < 1)) ? 1 : v; + } + + private static Semaphore sem(Kind k) { + return (k == Kind.SIMILARITY) ? SIMILARITY_SEM : GENERAL_SEM; + } + + private static String label(Kind k) { + return (k == Kind.SIMILARITY) ? "similarity" : "general"; + } + + /** + * Non-blocking acquire. Returns an {@link Permit} to be closed in a try-with-resources / finally, + * or {@code null} when the pool is at capacity (the caller should respond 429 with Retry-After). + */ + public static Permit tryAcquire(Kind k) { + if (!sem(k).tryAcquire()) { + metric(() -> { if (REJECTED != null) REJECTED.labels(label(k)).inc(); }); + return null; + } + // permit is held; metric updates are best-effort and must never lose the permit + metric(() -> { if (INFLIGHT != null) INFLIGHT.labels(label(k)).inc(); }); + return new Permit(k); + } + + private static void metric(Runnable r) { + try { r.run(); } catch (RuntimeException ignore) { /* metrics are best-effort */ } + } + + /** Idempotent permit handle — releasing exactly once, safe to close more than once. */ + public static final class Permit implements AutoCloseable { + private final Kind kind; + private boolean released = false; + + private Permit(Kind k) { + this.kind = k; + } + + @Override public void close() { + if (released) return; + released = true; + sem(kind).release(); + metric(() -> { if (INFLIGHT != null) INFLIGHT.labels(label(kind)).dec(); }); + } + } + + // --- test seams --- + static int availablePermits(Kind k) { + return sem(k).availablePermits(); + } +} diff --git a/src/main/webapp/WEB-INF/web.xml b/src/main/webapp/WEB-INF/web.xml index 9ac50f0fcc..04cdffc958 100644 --- a/src/main/webapp/WEB-INF/web.xml +++ b/src/main/webapp/WEB-INF/web.xml @@ -112,6 +112,7 @@ # No Bearer => falls through to SearchApi (which self-401s if no session). Filter handles auth. /api/v3/search/** = tokenAuthSearch /api/v3/media/resolve = tokenAuthSearch + /api/v3/similarity = tokenAuthSearch /api/v3/agent-skill = anon /api/v3/agent-skill/** = anon @@ -506,6 +507,16 @@ /api/v3/search/* + + SimilarityApi + org.ecocean.api.SimilarityApi + + + + SimilarityApi + /api/v3/similarity + + SiteSettings org.ecocean.api.SiteSettings diff --git a/src/test/java/org/ecocean/api/MediaResolveApiTest.java b/src/test/java/org/ecocean/api/MediaResolveApiTest.java index a9a72ca3ee..d4d601b737 100644 --- a/src/test/java/org/ecocean/api/MediaResolveApiTest.java +++ b/src/test/java/org/ecocean/api/MediaResolveApiTest.java @@ -656,4 +656,35 @@ private Annotation fullAnnotation(String id, String encId, String indId) throws } assertEquals(0, new JSONArray(out.toString()).length(), "fully out-of-bounds bbox -> omit"); } + + @Test void tokenRequest_whenGeneralPoolExhausted_429() throws Exception { + // drain the GENERAL pool so the token request cannot acquire a permit -> 429 (live users are + // protected from a high-volume agent) + java.util.List held = new java.util.ArrayList<>(); + TokenApiConcurrency.Permit p; + while ((p = TokenApiConcurrency.tryAcquire(TokenApiConcurrency.Kind.GENERAL)) != null) held.add(p); + try { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + StringWriter out = new StringWriter(); + when(resp.getWriter()).thenReturn(new PrintWriter(out)); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = mockConstruction(Shepherd.class, (m, c) -> { + doNothing().when(m).beginDBTransaction(); + doNothing().when(m).setAction(anyString()); + doNothing().when(m).rollbackAndClose(); + User u = mockUser("admin"); + when(m.getUser(any(HttpServletRequest.class))).thenReturn(u); + when(u.isAdmin(m)).thenReturn(true); + })) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationIds", new JSONArray().put("ann-A"))); + new MediaResolveApi().doPostForTest(req, resp); + } + verify(resp).setStatus(429); + verify(resp).setHeader(eq("Retry-After"), anyString()); + } finally { + for (TokenApiConcurrency.Permit h : held) h.close(); + } + } } diff --git a/src/test/java/org/ecocean/api/SimilarityApiTest.java b/src/test/java/org/ecocean/api/SimilarityApiTest.java new file mode 100644 index 0000000000..927838b041 --- /dev/null +++ b/src/test/java/org/ecocean/api/SimilarityApiTest.java @@ -0,0 +1,270 @@ +package org.ecocean.api; + +import static org.junit.jupiter.api.Assertions.*; +import static org.mockito.ArgumentMatchers.*; +import static org.mockito.Mockito.*; + +import java.io.PrintWriter; +import java.io.StringWriter; +import java.util.ArrayList; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Set; +import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; +import org.ecocean.Annotation; +import org.ecocean.Embedding; +import org.ecocean.Encounter; +import org.ecocean.OpenSearch; +import org.ecocean.User; +import org.ecocean.security.WildbookTokenAuthenticationFilter; +import org.ecocean.servlet.ServletUtilities; +import org.ecocean.shepherd.core.Shepherd; +import org.json.JSONArray; +import org.json.JSONObject; +import org.junit.jupiter.api.Test; +import org.mockito.MockedConstruction; +import org.mockito.MockedStatic; + +class SimilarityApiTest { + + private HttpServletRequest tokenRequest() { + HttpServletRequest req = mock(HttpServletRequest.class); + when(req.getMethod()).thenReturn("POST"); + when(req.getAttribute(WildbookTokenAuthenticationFilter.TOKEN_AUTH_ATTR)).thenReturn(Boolean.TRUE); + when(req.getAttribute(WildbookTokenAuthenticationFilter.TOKEN_CONTEXT_ATTR)).thenReturn("context0"); + return req; + } + + private User mockUser(String id) { + User u = mock(User.class); + when(u.getId()).thenReturn(id); + return u; + } + + private Annotation annWithEmbedding(String id, String viewpoint, String enc) { + Embedding e = mock(Embedding.class); + when(e.getMethod()).thenReturn("miewid"); + when(e.getMethodVersion()).thenReturn("msv4.1"); + when(e.vectorToFloatArray()).thenReturn(new float[] {0.1f, 0.2f, 0.3f}); + Set embs = new LinkedHashSet<>(); + embs.add(e); + Annotation ann = mock(Annotation.class); + when(ann.getId()).thenReturn(id); + when(ann.getViewpoint()).thenReturn(viewpoint); + when(ann.getEmbeddings()).thenReturn(embs); + Encounter en = mock(Encounter.class); + when(en.getId()).thenReturn(enc); + when(ann.findEncounter(any(Shepherd.class))).thenReturn(en); + return ann; + } + + private MockedConstruction shepherd(User user, boolean admin, Annotation ann) { + return mockConstruction(Shepherd.class, (m, c) -> { + doNothing().when(m).beginDBTransaction(); + doNothing().when(m).setAction(anyString()); + doNothing().when(m).rollbackAndClose(); + when(m.getUser(any(HttpServletRequest.class))).thenReturn(user); + when(user.isAdmin(m)).thenReturn(admin); + if (ann != null) when(m.getAnnotation(ann.getId())).thenReturn(ann); + }); + } + + private JSONObject knnResult(double score, String neighborId, String vp) { + JSONObject src = new JSONObject().put("id", neighborId).put("viewpoint", vp) + .put("encounterId", "e-" + neighborId).put("individualId", "i-" + neighborId); + JSONObject hit = new JSONObject().put("_score", score).put("_id", neighborId).put("_source", src); + return new JSONObject().put("hits", new JSONObject().put("hits", new JSONArray().put(hit))); + } + + @Test void nonToken_401() throws Exception { + HttpServletRequest req = mock(HttpServletRequest.class); + when(req.getMethod()).thenReturn("POST"); + when(req.getAttribute(WildbookTokenAuthenticationFilter.TOKEN_AUTH_ATTR)).thenReturn(Boolean.FALSE); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + new SimilarityApi().doPostForTest(req, resp); + verify(resp).setStatus(401); + } + + @Test void missingAnnotationId_400() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = mockConstruction(Shepherd.class)) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("k", 5)); + new SimilarityApi().doPostForTest(req, resp); + assertTrue(sh.constructed().isEmpty(), "no Shepherd before body validation"); + } + verify(resp).setStatus(400); + } + + @Test void badK_400() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + try (MockedStatic su = mockStatic(ServletUtilities.class)) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "a").put("k", 9999)); + new SimilarityApi().doPostForTest(req, resp); + } + verify(resp).setStatus(400); + } + + @Test void admin_happyPath_returnsNeighborsWithScore_noVectors_filteredAnnShape() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + StringWriter out = new StringWriter(); + when(resp.getWriter()).thenReturn(new PrintWriter(out)); + Annotation ann = annWithEmbedding("ann-A", "left", "enc-A"); + User admin = mockUser("admin"); + final JSONObject[] sentBody = {null}; + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = shepherd(admin, true, ann); + MockedConstruction os = mockConstruction(OpenSearch.class, (m, c) -> + when(m.queryPitPrivate(eq("annotation"), any(JSONObject.class), anyInt(), anyInt(), + any(), any())).thenAnswer(inv -> { + sentBody[0] = inv.getArgument(1); + return knnResult(0.91, "n1", "left"); + }))) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "ann-A").put("k", 5)); + new SimilarityApi().doPostForTest(req, resp); + } + verify(resp).setStatus(200); + JSONObject body = new JSONObject(out.toString()); + assertEquals("ann-A", body.getJSONObject("query").getString("annotationId")); + assertEquals("left", body.getJSONObject("query").getString("viewpoint")); + assertEquals("msv4.1", body.getJSONObject("query").getString("methodVersion")); + JSONArray nb = body.getJSONArray("neighbors"); + assertEquals(1, nb.length()); + assertEquals("n1", nb.getJSONObject(0).getString("id")); + assertEquals(0.91, nb.getJSONObject(0).getDouble("score"), 0.0001); + assertFalse(nb.getJSONObject(0).has("embeddings"), "neighbors never carry vectors"); + // the kNN body must filter INSIDE knn.filter (filtered ANN), with viewpoint+method+version and self-exclusion + JSONObject knn = sentBody[0].getJSONObject("query").getJSONObject("nested") + .getJSONObject("query").getJSONObject("knn").getJSONObject("embeddings.vector"); + JSONObject fbool = knn.getJSONObject("filter").getJSONObject("bool"); + String filt = fbool.getJSONArray("filter").toString(); + assertTrue(filt.contains("viewpoint") && filt.contains("left"), "filters to the query viewpoint"); + assertTrue(filt.contains("embeddings.methodVersion") && filt.contains("msv4.1"), + "filters to the query method version"); + assertTrue(fbool.getJSONArray("must_not").toString().contains("ann-A"), + "excludes the query annotation itself"); + assertTrue(sentBody[0].getJSONObject("_source").getJSONArray("excludes").toString() + .contains("embeddings"), "embeddings excluded from _source"); + } + + @Test void neighborEmbeddingsInSource_areStripped() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + StringWriter out = new StringWriter(); + when(resp.getWriter()).thenReturn(new PrintWriter(out)); + Annotation ann = annWithEmbedding("ann-A", "left", "enc-A"); + User admin = mockUser("admin"); + // a hit whose _source DOES contain an embeddings vector — must be stripped before output + JSONObject src = new JSONObject().put("id", "n1").put("viewpoint", "left") + .put("embeddings", new JSONArray().put(new JSONObject() + .put("methodVersion", "msv4.1").put("vector", new JSONArray().put(0.1).put(0.2)))); + JSONObject hit = new JSONObject().put("_score", 0.8).put("_source", src); + JSONObject res = new JSONObject().put("hits", new JSONObject().put("hits", + new JSONArray().put(hit))); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = shepherd(admin, true, ann); + MockedConstruction os = mockConstruction(OpenSearch.class, (m, c) -> + when(m.queryPitPrivate(eq("annotation"), any(JSONObject.class), anyInt(), anyInt(), + any(), any())).thenReturn(res))) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "ann-A")); + new SimilarityApi().doPostForTest(req, resp); + } + JSONObject nb = new JSONObject(out.toString()).getJSONArray("neighbors").getJSONObject(0); + assertFalse(nb.has("embeddings"), "embeddings in _source must be stripped from the response"); + } + + @Test void ambiguousEmbedding_400() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + // two embeddings both matching method=miewid; no methodVersion given -> ambiguous -> 400 + Embedding e1 = mock(Embedding.class); + when(e1.getMethod()).thenReturn("miewid"); when(e1.getMethodVersion()).thenReturn("msv4.1"); + Embedding e2 = mock(Embedding.class); + when(e2.getMethod()).thenReturn("miewid"); when(e2.getMethodVersion()).thenReturn("msv3"); + Set embs = new LinkedHashSet<>(); embs.add(e1); embs.add(e2); + Annotation ann = mock(Annotation.class); + when(ann.getId()).thenReturn("ann-amb"); + when(ann.getViewpoint()).thenReturn("left"); + when(ann.getEmbeddings()).thenReturn(embs); + User admin = mockUser("admin"); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = shepherd(admin, true, ann); + MockedConstruction os = mockConstruction(OpenSearch.class)) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "ann-amb")); // no methodVersion + new SimilarityApi().doPostForTest(req, resp); + } + verify(resp).setStatus(400); + } + + @Test void noViewpoint_400() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + Annotation ann = annWithEmbedding("ann-nv", null, "enc"); // null viewpoint + User admin = mockUser("admin"); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = shepherd(admin, true, ann); + MockedConstruction os = mockConstruction(OpenSearch.class)) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "ann-nv")); + new SimilarityApi().doPostForTest(req, resp); + } + verify(resp).setStatus(400); + } + + @Test void nonAdmin_queryAnnotationNotVisible_404_noOracle() throws Exception { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + User viewer = mockUser("viewer"); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = shepherd(viewer, false, null); + MockedConstruction os = mockConstruction(OpenSearch.class, (m, c) -> + // the visibility eligibility query returns no hits -> not visible + when(m.queryPitPrivate(eq("annotation"), any(JSONObject.class), anyInt(), anyInt(), + any(), any())).thenReturn(new JSONObject().put("hits", + new JSONObject().put("hits", new JSONArray()))))) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "hidden-or-missing")); + new SimilarityApi().doPostForTest(req, resp); + verify(sh.constructed().get(0), never()).getAnnotation("hidden-or-missing"); + } + verify(resp).setStatus(404); + } + + @Test void poolExhausted_429_withRetryAfter_noShepherd() throws Exception { + // drain the SIMILARITY pool so the request can't acquire a permit + List held = new ArrayList<>(); + TokenApiConcurrency.Permit p; + while ((p = TokenApiConcurrency.tryAcquire(TokenApiConcurrency.Kind.SIMILARITY)) != null) held.add(p); + try { + HttpServletRequest req = tokenRequest(); + HttpServletResponse resp = mock(HttpServletResponse.class); + when(resp.getWriter()).thenReturn(new PrintWriter(new StringWriter())); + try (MockedStatic su = mockStatic(ServletUtilities.class); + MockedConstruction sh = mockConstruction(Shepherd.class)) { + su.when(() -> ServletUtilities.jsonFromHttpServletRequest(any())) + .thenReturn(new JSONObject().put("annotationId", "ann-A").put("k", 5)); + new SimilarityApi().doPostForTest(req, resp); + assertTrue(sh.constructed().isEmpty(), "no Shepherd/DB work when the pool is at capacity"); + } + verify(resp).setStatus(429); + verify(resp).setHeader(eq("Retry-After"), anyString()); + } finally { + for (TokenApiConcurrency.Permit h : held) h.close(); + } + } +} diff --git a/src/test/java/org/ecocean/api/TokenApiConcurrencyTest.java b/src/test/java/org/ecocean/api/TokenApiConcurrencyTest.java new file mode 100644 index 0000000000..fc79a9619d --- /dev/null +++ b/src/test/java/org/ecocean/api/TokenApiConcurrencyTest.java @@ -0,0 +1,67 @@ +package org.ecocean.api; + +import static org.junit.jupiter.api.Assertions.*; + +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Test; + +/** + * The per-instance concurrency guard: a bounded number of token requests run at once (excess gets a + * null permit -> caller 429s), permits are restored on close, the two pools are independent, and + * double-close is safe (no over-release). + */ +class TokenApiConcurrencyTest { + + private List drain(TokenApiConcurrency.Kind k) { + List held = new ArrayList<>(); + TokenApiConcurrency.Permit p; + while ((p = TokenApiConcurrency.tryAcquire(k)) != null) held.add(p); + return held; + } + + @Test void exhaustsThenRecovers() { + TokenApiConcurrency.Kind K = TokenApiConcurrency.Kind.SIMILARITY; + List held = drain(K); + try { + assertTrue(held.size() >= 1, "pool has at least one permit"); + assertEquals(0, TokenApiConcurrency.availablePermits(K), "pool is fully drained"); + assertNull(TokenApiConcurrency.tryAcquire(K), "no permit when the pool is at capacity"); + + held.get(0).close(); // release one + TokenApiConcurrency.Permit again = TokenApiConcurrency.tryAcquire(K); + assertNotNull(again, "a permit is available again after a release"); + again.close(); + } finally { + for (TokenApiConcurrency.Permit p : held) p.close(); + } + // all returned + int max = TokenApiConcurrency.availablePermits(K); + assertTrue(max >= 1, "permits fully restored after closing"); + } + + @Test void poolsAreIndependent() { + List sim = drain(TokenApiConcurrency.Kind.SIMILARITY); + try { + assertEquals(0, TokenApiConcurrency.availablePermits(TokenApiConcurrency.Kind.SIMILARITY)); + TokenApiConcurrency.Permit g = + TokenApiConcurrency.tryAcquire(TokenApiConcurrency.Kind.GENERAL); + assertNotNull(g, "exhausting SIMILARITY must not block GENERAL"); + g.close(); + } finally { + for (TokenApiConcurrency.Permit p : sim) p.close(); + } + } + + @Test void doubleCloseIsSafe() { + TokenApiConcurrency.Kind K = TokenApiConcurrency.Kind.GENERAL; + int before = TokenApiConcurrency.availablePermits(K); + TokenApiConcurrency.Permit p = TokenApiConcurrency.tryAcquire(K); + assertNotNull(p); + assertEquals(before - 1, TokenApiConcurrency.availablePermits(K)); + p.close(); + p.close(); // must NOT over-release + assertEquals(before, TokenApiConcurrency.availablePermits(K), + "double close releases exactly one permit"); + } +}