From 3a3f43b380cbfdd948b64c9e41833bf7f4f77734 Mon Sep 17 00:00:00 2001 From: Serena Shen <83137666+ChaoyangS@users.noreply.github.com> Date: Sun, 14 Jun 2026 13:21:18 -0400 Subject: [PATCH 1/2] feat(core): add operation-level sample rates to the usage plugin --- .changeset/operation-level-sample-rates.md | 19 +++ .../core/demo/operation-sample-rates.ts | 105 +++++++++++++ .../libraries/core/src/client/sampling.ts | 55 ++++++- packages/libraries/core/src/client/types.ts | 47 +++++- packages/libraries/core/src/client/usage.ts | 9 +- .../libraries/core/tests/sampling.spec.ts | 143 ++++++++++++++++++ packages/libraries/core/tests/usage.spec.ts | 78 ++++++++++ 7 files changed, 452 insertions(+), 4 deletions(-) create mode 100644 .changeset/operation-level-sample-rates.md create mode 100644 packages/libraries/core/demo/operation-sample-rates.ts create mode 100644 packages/libraries/core/tests/sampling.spec.ts diff --git a/.changeset/operation-level-sample-rates.md b/.changeset/operation-level-sample-rates.md new file mode 100644 index 00000000000..09c2f618daa --- /dev/null +++ b/.changeset/operation-level-sample-rates.md @@ -0,0 +1,19 @@ +--- +'@graphql-hive/core': minor +--- + +Add operation-level sample rates to the usage plugin via the new `sampleRates` option. + +Each rule matches operations by exact `name` or by `regex` (tested against the operation name) and applies its own `sampleRate`. The first matching rule wins; operations that don't match any rule fall back to the global `sampleRate` (default `1.0`). This lets you sample known high-volume operations aggressively while keeping 100% visibility into low-volume operations, instead of choosing a single global rate that either overpays on quota or misses rare operations. + +`exclude` still takes precedence (excluded operations are never reported), and `sampler`, when provided, continues to override both `sampleRate` and `sampleRates`. + +```ts +usage: { + sampleRate: 1.0, // fallback for unmatched operations + sampleRates: [ + { name: 'SampledQuery', sampleRate: 0.1 }, + { regex: /^HighVolume/, sampleRate: 0.1 }, + ], +} +``` diff --git a/packages/libraries/core/demo/operation-sample-rates.ts b/packages/libraries/core/demo/operation-sample-rates.ts new file mode 100644 index 00000000000..aaae3d4b0b0 --- /dev/null +++ b/packages/libraries/core/demo/operation-sample-rates.ts @@ -0,0 +1,105 @@ +/** + * Demo / reproduction for operation-level sample rates. + * + * Scenario from the feature request: a system mixes a few very high-volume + * operations with many low-volume ones. Global sampling forces a bad trade-off: + * - sample everything at 100% -> the high-volume ops blow the usage quota + * - sample everything at e.g. 5% -> the low-volume ops are missed entirely + * + * Operation-level sample rates let us sample the known high-volume operations + * aggressively while keeping 100% visibility into low-volume operations. + * + * Run from the repo root: + * pnpm --filter @graphql-hive/core exec tsx demo/operation-sample-rates.ts + * or from packages/libraries/core: + * pnpm exec tsx demo/operation-sample-rates.ts + */ +import { parse } from 'graphql'; +import { operationSampling, randomSampling } from '../src/client/sampling.js'; +import type { SamplingContext } from '../src/client/types.js'; + +const document = parse(/* GraphQL */ ` + query { + __typename + } +`); + +function ctx(operationName: string): SamplingContext { + return { operationName, document, variableValues: null, contextValue: undefined }; +} + +// A simulated traffic mix: operation name -> number of executions in the window. +const traffic: Array<{ operationName: string; volume: number }> = [ + { operationName: 'HighVolumeFeed', volume: 500_000 }, + { operationName: 'HighVolumePing', volume: 250_000 }, + { operationName: 'CheckoutMutation', volume: 1_200 }, + { operationName: 'AdminAuditReport', volume: 40 }, + { operationName: 'RareMigrationQuery', volume: 8 }, +]; + +const totalOps = traffic.reduce((sum, t) => sum + t.volume, 0); + +function simulate(shouldInclude: (c: SamplingContext) => boolean) { + const sent = new Map(); + for (const { operationName, volume } of traffic) { + let count = 0; + for (let i = 0; i < volume; i++) { + if (shouldInclude(ctx(operationName))) count++; + } + sent.set(operationName, count); + } + return sent; +} + +function pct(part: number, whole: number) { + return whole === 0 ? '0%' : `${((part / whole) * 100).toFixed(1)}%`; +} + +function printResult(title: string, sent: Map) { + console.log(`\n=== ${title} ===`); + let totalSent = 0; + for (const { operationName, volume } of traffic) { + const s = sent.get(operationName) ?? 0; + totalSent += s; + const visibility = s > 0 ? 'visible' : 'MISSED'; + console.log( + ` ${operationName.padEnd(20)} volume=${String(volume).padStart(7)} ` + + `sent=${String(s).padStart(7)} (${pct(s, volume).padStart(6)}) ${visibility}`, + ); + } + console.log( + ` ${'TOTAL'.padEnd(20)} volume=${String(totalOps).padStart(7)} ` + + `sent=${String(totalSent).padStart(7)} (${pct(totalSent, totalOps)} of quota)`, + ); + return totalSent; +} + +// --- Strategy A: global 100% sampling (full visibility, full quota cost) ----- +const globalFull = printResult('Global sampleRate = 1.0 (overpay)', simulate(randomSampling(1.0))); + +// --- Strategy B: global 5% sampling (cheap, but loses low-volume ops) -------- +const globalLow = printResult( + 'Global sampleRate = 0.05 (low-volume ops missed)', + simulate(randomSampling(0.05)), +); + +// --- Strategy C: operation-level sample rates (the new feature) -------------- +// Sample the high-volume ops at 1%, keep everything else at 100% via fallback. +const perOp = printResult( + 'Operation-level: /^HighVolume/ -> 1%, fallback 100%', + simulate( + operationSampling({ + rates: [{ regex: /^HighVolume/, sampleRate: 0.01 }], + fallbackSampleRate: 1.0, + }), + ), +); + +console.log('\n=== Summary ==='); +console.log(` Global 100%: ${globalFull} ops billed (${pct(globalFull, totalOps)})`); +console.log(` Global 5%: ${globalLow} ops billed (${pct(globalLow, totalOps)})`); +console.log(` Operation-level: ${perOp} ops billed (${pct(perOp, totalOps)})`); +console.log( + '\n -> Operation-level sampling slashes billed volume like the cheap global rate,', +); +console.log(' while still reporting 100% of the low-volume operations.'); diff --git a/packages/libraries/core/src/client/sampling.ts b/packages/libraries/core/src/client/sampling.ts index 5edb210d44c..694234b3f2e 100644 --- a/packages/libraries/core/src/client/sampling.ts +++ b/packages/libraries/core/src/client/sampling.ts @@ -1,4 +1,4 @@ -import type { SamplingContext } from './types.js'; +import type { OperationSampleRate, SamplingContext } from './types.js'; export function randomSampling(sampleRate: number) { if (sampleRate > 1 || sampleRate < 0) { @@ -10,6 +10,59 @@ export function randomSampling(sampleRate: number) { }; } +/** + * Builds a sampling function that applies operation-level sample rates. + * + * Each rule matches an operation by exact `name` or by `regex` (tested against the + * operation name). The first matching rule wins and its `sampleRate` is used. + * Operations that don't match any rule fall back to `fallbackSampleRate`. + * + * Rules are validated eagerly so misconfiguration fails fast at setup time. + */ +export function operationSampling(config: { + rates: OperationSampleRate[]; + fallbackSampleRate: number; +}) { + if (config.fallbackSampleRate > 1 || config.fallbackSampleRate < 0) { + throw new Error( + `Expected usage.sampleRate to be 0 <= x <= 1, received ${config.fallbackSampleRate}`, + ); + } + + const rules = config.rates.map(rate => { + const hasName = typeof rate.name === 'string'; + const hasRegex = rate.regex instanceof RegExp; + + if (hasName === hasRegex) { + throw new Error( + 'Expected usage.sampleRates entry to define exactly one of "name" or "regex".', + ); + } + + if (rate.sampleRate > 1 || rate.sampleRate < 0) { + throw new Error( + `Expected usage.sampleRates sampleRate to be 0 <= x <= 1, received ${rate.sampleRate}`, + ); + } + + const matches = hasRegex + ? (operationName: string) => rate.regex!.test(operationName) + : (operationName: string) => operationName === rate.name; + + return { matches, sampleRate: rate.sampleRate }; + }); + + return function shouldInclude(context: SamplingContext): boolean { + for (const rule of rules) { + if (rule.matches(context.operationName)) { + return Math.random() <= rule.sampleRate; + } + } + + return Math.random() <= config.fallbackSampleRate; + }; +} + export function dynamicSampling(sampler: (context: SamplingContext) => number | boolean) { return function shouldInclude(context: SamplingContext): boolean { let sampleRate = sampler(context); diff --git a/packages/libraries/core/src/client/types.ts b/packages/libraries/core/src/client/types.ts index 4d154016fef..e4f3fa1eb22 100644 --- a/packages/libraries/core/src/client/types.ts +++ b/packages/libraries/core/src/client/types.ts @@ -132,10 +132,31 @@ export interface HiveUsagePluginOptions { * Default: 1.0 */ sampleRate?: number; + /** + * Operation-level sample rates. + * + * Each rule matches operations by exact `name` or by `regex` (tested against the + * operation name) and applies its own `sampleRate`. The first matching rule wins. + * Operations that don't match any rule fall back to the global `sampleRate` + * (or 1.0 when that is not set). + * + * This is useful when a system has a mix of high-volume and low-volume operations: + * high-volume operations can be sampled aggressively while low-volume operations + * are reported at (or close to) 100%. + * + * Ignored when `sampler` is defined. + * + * @example + * sampleRates: [ + * { name: 'SampledQuery', sampleRate: 0.1 }, + * { regex: /^Sampled/, sampleRate: 0.1 }, + * ] + */ + sampleRates?: OperationSampleRate[]; /** * Compute sample rate dynamically. * - * If `sampler` is defined, `sampleRate` is ignored. + * If `sampler` is defined, `sampleRate` and `sampleRates` are ignored. * * @returns A sample rate between 0 and 1. * 0.0 = 0% chance of being sent @@ -152,6 +173,30 @@ export interface HiveUsagePluginOptions { processVariables?: boolean; } +/** + * A sample rate that applies to a subset of operations, matched by name or regular expression. + */ +export interface OperationSampleRate { + /** + * Match operations by their exact operation name. + * + * Exactly one of `name` or `regex` must be provided. + */ + name?: string; + /** + * Match operations whose name matches this regular expression. + * + * Exactly one of `name` or `regex` must be provided. + */ + regex?: RegExp; + /** + * Sample rate applied to operations matched by this rule. + * 0.0 = 0% chance of being sent + * 1.0 = 100% chance of being sent. + */ + sampleRate: number; +} + export interface SamplingContext extends Pick { operationName: string; diff --git a/packages/libraries/core/src/client/usage.ts b/packages/libraries/core/src/client/usage.ts index 275d50e0f2c..fb8c678861a 100644 --- a/packages/libraries/core/src/client/usage.ts +++ b/packages/libraries/core/src/client/usage.ts @@ -11,7 +11,7 @@ import { normalizeOperation } from '../normalize/operation.js'; import { version } from '../version.js'; import { createAgent } from './agent.js'; import { collectSchemaCoordinates } from './collect-schema-coordinates.js'; -import { dynamicSampling, randomSampling } from './sampling.js'; +import { dynamicSampling, operationSampling, randomSampling } from './sampling.js'; import type { AbortAction, ClientInfo, @@ -191,7 +191,12 @@ export function createUsage(pluginOptions: HiveInternalPluginOptions): UsageColl const shouldInclude = options.sampler && typeof options.sampler === 'function' ? dynamicSampling(options.sampler) - : randomSampling(options.sampleRate ?? 1.0); + : options.sampleRates?.length + ? operationSampling({ + rates: options.sampleRates, + fallbackSampleRate: options.sampleRate ?? 1.0, + }) + : randomSampling(options.sampleRate ?? 1.0); const collectRequest: UsageCollector['collectRequest'] = args => { let providedOperationName: string | undefined = undefined; diff --git a/packages/libraries/core/tests/sampling.spec.ts b/packages/libraries/core/tests/sampling.spec.ts new file mode 100644 index 00000000000..160d089c468 --- /dev/null +++ b/packages/libraries/core/tests/sampling.spec.ts @@ -0,0 +1,143 @@ +import { parse } from 'graphql'; +import { operationSampling, randomSampling } from '../src/client/sampling'; +import type { SamplingContext } from '../src/client/types'; + +const document = parse(/* GraphQL */ ` + query { + __typename + } +`); + +function context(operationName: string): SamplingContext { + return { + operationName, + document, + variableValues: null, + contextValue: undefined, + }; +} + +afterEach(() => { + vi.restoreAllMocks(); +}); + +describe('operationSampling', () => { + test('uses the rule sample rate for an exact name match', () => { + const random = vi.spyOn(Math, 'random').mockReturnValue(0.5); + const shouldInclude = operationSampling({ + rates: [{ name: 'HighVolume', sampleRate: 0.1 }], + fallbackSampleRate: 1, + }); + + // 0.5 <= 0.1 -> excluded + expect(shouldInclude(context('HighVolume'))).toBe(false); + + random.mockReturnValue(0.05); + // 0.05 <= 0.1 -> included + expect(shouldInclude(context('HighVolume'))).toBe(true); + }); + + test('uses the rule sample rate for a regex match', () => { + vi.spyOn(Math, 'random').mockReturnValue(0.5); + const shouldInclude = operationSampling({ + rates: [{ regex: /^Sampled/, sampleRate: 0.1 }], + fallbackSampleRate: 1, + }); + + expect(shouldInclude(context('SampledQuery'))).toBe(false); + // does not match the regex -> falls back to 1.0 -> always included + expect(shouldInclude(context('OtherQuery'))).toBe(true); + }); + + test('falls back to the global sample rate when no rule matches', () => { + vi.spyOn(Math, 'random').mockReturnValue(0.5); + const shouldInclude = operationSampling({ + rates: [{ name: 'HighVolume', sampleRate: 0.1 }], + fallbackSampleRate: 0, + }); + + // No match -> fallback 0 -> never included (low-volume ops would normally be 1.0 here) + expect(shouldInclude(context('LowVolume'))).toBe(false); + }); + + test('first matching rule wins', () => { + vi.spyOn(Math, 'random').mockReturnValue(0); // 0 <= any rate >= 0 -> included unless rate is 0 + const shouldInclude = operationSampling({ + rates: [ + { regex: /Query$/, sampleRate: 1 }, + { name: 'GetUserQuery', sampleRate: 0 }, + ], + fallbackSampleRate: 0, + }); + + // Matches the first rule (regex) even though the second exact rule also matches. + expect(shouldInclude(context('GetUserQuery'))).toBe(true); + }); + + test('keeps the high-volume vs low-volume scenario from the feature request', () => { + // Reports ~10% of high-volume ops and 100% of everything else. + const shouldInclude = operationSampling({ + rates: [{ regex: /^HighVolume/, sampleRate: 0.1 }], + fallbackSampleRate: 1, + }); + + const random = vi.spyOn(Math, 'random'); + + // High-volume op: below threshold -> sent + random.mockReturnValue(0.05); + expect(shouldInclude(context('HighVolumeQuery'))).toBe(true); + // High-volume op: above threshold -> dropped + random.mockReturnValue(0.9); + expect(shouldInclude(context('HighVolumeQuery'))).toBe(false); + + // Low-volume op: always sent regardless of the random draw + random.mockReturnValue(0.99); + expect(shouldInclude(context('RareReport'))).toBe(true); + }); + + describe('validation', () => { + test('throws when a rule defines neither name nor regex', () => { + expect(() => + operationSampling({ + rates: [{ sampleRate: 0.1 } as any], + fallbackSampleRate: 1, + }), + ).toThrow('exactly one of "name" or "regex"'); + }); + + test('throws when a rule defines both name and regex', () => { + expect(() => + operationSampling({ + rates: [{ name: 'Foo', regex: /Foo/, sampleRate: 0.1 }], + fallbackSampleRate: 1, + }), + ).toThrow('exactly one of "name" or "regex"'); + }); + + test('throws when a rule sample rate is out of range', () => { + expect(() => + operationSampling({ + rates: [{ name: 'Foo', sampleRate: 1.5 }], + fallbackSampleRate: 1, + }), + ).toThrow('0 <= x <= 1'); + }); + + test('throws when the fallback sample rate is out of range', () => { + expect(() => + operationSampling({ + rates: [{ name: 'Foo', sampleRate: 0.1 }], + fallbackSampleRate: -1, + }), + ).toThrow('0 <= x <= 1'); + }); + }); +}); + +describe('randomSampling (regression)', () => { + test('still samples against the provided rate', () => { + vi.spyOn(Math, 'random').mockReturnValue(0.5); + expect(randomSampling(0.4)()).toBe(false); + expect(randomSampling(0.6)()).toBe(true); + }); +}); diff --git a/packages/libraries/core/tests/usage.spec.ts b/packages/libraries/core/tests/usage.spec.ts index dd03a80be10..681a8ec9037 100644 --- a/packages/libraries/core/tests/usage.spec.ts +++ b/packages/libraries/core/tests/usage.spec.ts @@ -702,6 +702,84 @@ test('should not send excluded operation name data to Hive', async () => { expect(operation.execution.ok).toBe(true); }); +test('should apply operation-level sample rates with global fallback', async () => { + const logger = createHiveTestingLogger(); + + // Deterministic random draw of 0.5: + // - high-volume ops (rate 0.1): 0.5 <= 0.1 === false -> dropped + // - everything else (fallback rate 1.0): 0.5 <= 1.0 === true -> kept + vi.spyOn(Math, 'random').mockReturnValue(0.5); + + const token = 'Token'; + + let report: Report = { + size: 0, + map: {}, + operations: [], + }; + const http = nock('http://localhost') + .post('/200') + .once() + .reply((_, _body) => { + report = _body as any; + return [200]; + }); + + const hive = createHive({ + enabled: true, + debug: true, + agent: { + timeout: 500, + maxRetries: 0, + sendInterval: 10, + logger, + }, + token, + selfHosting: { + graphqlEndpoint: 'http://localhost/graphql', + applicationUrl: 'http://localhost/', + usageEndpoint: 'http://localhost/200', + }, + usage: { + // global fallback stays at 100% + sampleRate: 1, + sampleRates: [ + // exact-name rule + { name: 'deleteProject', sampleRate: 0.1 }, + // regex rule + { regex: /^HighVolume/, sampleRate: 0.1 }, + ], + }, + }); + + const collect = hive.collectUsage(); + await waitFor(20); + await Promise.all([ + // matched by exact name rule (0.1) -> dropped + collect({ schema, document: op, operationName: 'deleteProject' }, {}), + // matched by regex rule (0.1) -> dropped + collect({ schema, document: op2, operationName: 'HighVolumeProject' }, {}), + // no rule matches -> fallback 1.0 -> kept + collect({ schema, document: op2, operationName: 'getProject' }, {}), + ]); + await hive.dispose(); + await waitFor(50); + http.done(); + + // Only the low-volume operation (getProject) survives sampling. + expect(report.size).toEqual(1); + expect(Object.keys(report.map)).toHaveLength(1); + + const key = Object.keys(report.map)[0]; + const record = report.map[key]; + expect(record.operation).toMatch('query getProject'); + expect(record.operationName).toMatch('getProject'); + + const operations = report.operations; + expect(operations).toHaveLength(1); + expect(operations?.[0].operationMapKey).toEqual(key); +}); + test('retry on non-200', async () => { const logger = createHiveTestingLogger(); From d6cc90b1e4bad4eee3578a670a96cdcd25442551 Mon Sep 17 00:00:00 2001 From: Serena Shen <83137666+ChaoyangS@users.noreply.github.com> Date: Sun, 28 Jun 2026 22:09:57 -0400 Subject: [PATCH 2/2] refactor(core): drop non-null assertion in operationSampling via type narrowing --- packages/libraries/core/src/client/sampling.ts | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/packages/libraries/core/src/client/sampling.ts b/packages/libraries/core/src/client/sampling.ts index 694234b3f2e..8867e6dc1b9 100644 --- a/packages/libraries/core/src/client/sampling.ts +++ b/packages/libraries/core/src/client/sampling.ts @@ -30,8 +30,9 @@ export function operationSampling(config: { } const rules = config.rates.map(rate => { - const hasName = typeof rate.name === 'string'; - const hasRegex = rate.regex instanceof RegExp; + const { name, regex, sampleRate } = rate; + const hasName = typeof name === 'string'; + const hasRegex = regex instanceof RegExp; if (hasName === hasRegex) { throw new Error( @@ -39,17 +40,18 @@ export function operationSampling(config: { ); } - if (rate.sampleRate > 1 || rate.sampleRate < 0) { + if (sampleRate > 1 || sampleRate < 0) { throw new Error( - `Expected usage.sampleRates sampleRate to be 0 <= x <= 1, received ${rate.sampleRate}`, + `Expected usage.sampleRates sampleRate to be 0 <= x <= 1, received ${sampleRate}`, ); } - const matches = hasRegex - ? (operationName: string) => rate.regex!.test(operationName) - : (operationName: string) => operationName === rate.name; + const matches = + regex instanceof RegExp + ? (operationName: string) => regex.test(operationName) + : (operationName: string) => operationName === name; - return { matches, sampleRate: rate.sampleRate }; + return { matches, sampleRate }; }); return function shouldInclude(context: SamplingContext): boolean {