From 7a28e4b81b101fa93353a116fae23e0c21e23a53 Mon Sep 17 00:00:00 2001 From: Espen Hovlandsdal Date: Mon, 24 Jun 2024 16:43:30 -0700 Subject: [PATCH 1/3] feat(shard): allow domain sharding --- src/data/listen.ts | 14 ++- src/http/browserMiddleware.ts | 4 +- src/http/domainSharding.ts | 96 ++++++++++++++++++++ src/http/nodeMiddleware.ts | 3 + src/http/requestOptions.ts | 1 + src/types.ts | 8 ++ test/domainSharding.test.ts | 163 ++++++++++++++++++++++++++++++++++ 7 files changed, 287 insertions(+), 2 deletions(-) create mode 100644 src/http/domainSharding.ts create mode 100644 test/domainSharding.test.ts diff --git a/src/data/listen.ts b/src/data/listen.ts index 52ffe86d7..ac59813b8 100644 --- a/src/data/listen.ts +++ b/src/data/listen.ts @@ -1,5 +1,6 @@ import {Observable} from 'rxjs' +import {domainSharder as sharder} from '../http/domainSharding' import type {ObservableSanityClient, SanityClient} from '../SanityClient' import type {Any, ListenEvent, ListenOptions, ListenParams, MutationEvent} from '../types' import defaults from '../util/defaults' @@ -64,7 +65,11 @@ export function _listen = Record>( const listenOpts = pick(options, possibleOptions) const qs = encodeQueryString({query, params, options: {tag, ...listenOpts}}) - const uri = `${url}${_getDataUrl(this, 'listen', qs)}` + let uri = `${url}${_getDataUrl(this, 'listen', qs)}` + if (this.config().useDomainSharding) { + uri = sharder.getShardedUrl(uri) + } + if (uri.length > MAX_URL_LENGTH) { return new Observable((observer) => observer.error(new Error('Query too large for listener'))) } @@ -91,6 +96,12 @@ export function _listen = Record>( // Once it is`true`, it will never be `false` again. let unsubscribed = false + // We're about to connect, and will reuse the same shard/bucket for every reconnect henceforth. + // This may seem inoptimal, but once connected we should just consider this as a "permanent" + // connection, since we'll automatically retry on failures/disconnects. Once we explicitly + // unsubsccribe, we can decrement the bucket and free up the shard. + sharder.incrementBucketForUrl(uri) + open() function onError() { @@ -187,6 +198,7 @@ export function _listen = Record>( stopped = true unsubscribe() unsubscribed = true + sharder.decrementBucketForUrl(uri) } return stop diff --git a/src/http/browserMiddleware.ts b/src/http/browserMiddleware.ts index 9859f0797..565a41f0e 100644 --- a/src/http/browserMiddleware.ts +++ b/src/http/browserMiddleware.ts @@ -1 +1,3 @@ -export default [] +import {domainSharder} from './domainSharding' + +export default [domainSharder.middleware] diff --git a/src/http/domainSharding.ts b/src/http/domainSharding.ts new file mode 100644 index 000000000..0b012538c --- /dev/null +++ b/src/http/domainSharding.ts @@ -0,0 +1,96 @@ +import type {Middleware, RequestOptions} from 'get-it' + +const UNSHARDED_URL_RE = /^https:\/\/([a-z0-9]+)\.api\.(sanity\..*)/ +const SHARDED_URL_RE = /^https:\/\/[a-z0-9]+\.api\.s(\d+)\.sanity\.(.*)/ + +/** + * Get a default sharding implementation where buckets are reused across instances. + * Helps prevent the case when multiple clients are instantiated, each having their + * own state of which buckets are least used. + */ +export const domainSharder = getDomainSharder() + +/** + * @internal + */ +export function getDomainSharder(initialBuckets?: number[]) { + const buckets: number[] = initialBuckets || new Array(10).fill(0, 0) + + function incrementBucketForUrl(url: string) { + const shard = getShardFromUrl(url) + if (shard !== null) { + buckets[shard]++ + } + } + + function decrementBucketForUrl(url: string) { + const shard = getShardFromUrl(url) + if (shard !== null) { + buckets[shard]-- + } + } + + function getShardedUrl(url: string): string { + const [isMatch, projectId, rest] = url.match(UNSHARDED_URL_RE) || [] + if (!isMatch) { + return url + } + + // Find index of bucket with fewest requests + const bucket = buckets.reduce( + (smallest, count, index) => (count < buckets[smallest] ? index : smallest), + 0, + ) + + return `https://${projectId}.api.s${bucket}.${rest}` + } + + function getShardFromUrl(url: string): number | null { + const [isMatch, shard] = url.match(SHARDED_URL_RE) || [] + return isMatch ? parseInt(shard, 10) : null + } + + const middleware = { + processOptions: (options: {useDomainSharding?: boolean; url: string}) => { + if (!useDomainSharding(options)) { + return options + } + + const url = getShardedUrl(options.url) + options.url = url + + return options + }, + + onRequest(req: { + options: Partial & {useDomainSharding?: boolean; url: string} + }) { + if (useDomainSharding(req.options)) { + incrementBucketForUrl(req.options.url) + } + return req + }, + + onResponse( + res, + context: {options: Partial & {useDomainSharding?: boolean; url: string}}, + ) { + if (useDomainSharding(context.options)) { + decrementBucketForUrl(context.options.url) + } + return res + }, + } satisfies Middleware + + return { + middleware, + incrementBucketForUrl, + decrementBucketForUrl, + getShardedUrl, + getBuckets: () => buckets, + } +} + +function useDomainSharding(options: RequestOptions | {useDomainSharding?: boolean}): boolean { + return 'useDomainSharding' in options && options.useDomainSharding === true +} diff --git a/src/http/nodeMiddleware.ts b/src/http/nodeMiddleware.ts index 2dff4f4f2..437c5f5e3 100644 --- a/src/http/nodeMiddleware.ts +++ b/src/http/nodeMiddleware.ts @@ -1,8 +1,11 @@ import {agent, debug, headers} from 'get-it/middleware' import {name, version} from '../../package.json' +import {domainSharder} from './domainSharding' const middleware = [ + domainSharder.middleware, + debug({verbose: true, namespace: 'sanity:client'}), headers({'User-Agent': `${name} ${version}`}), diff --git a/src/http/requestOptions.ts b/src/http/requestOptions.ts index fe31c1f5d..8d892c030 100644 --- a/src/http/requestOptions.ts +++ b/src/http/requestOptions.ts @@ -29,6 +29,7 @@ export function requestOptions(config: Any, overrides: Any = {}): Omit { + const isBrowser = typeof window !== 'undefined' && window.location && window.location.hostname + const isEdge = typeof EdgeRuntime === 'string' + let nock: typeof import('nock') = (() => { + throw new Error('Not supported in EdgeRuntime') + }) as any + if (!isEdge) { + const _nock = await import('nock') + nock = _nock.default + } + + const testClient = describe.each([ + [ + 'static config', + (conf?: ClientConfig) => + createClient({...clientConfig, useDomainSharding: true, ...(conf || {})}), + ], + [ + 'reconfigured config', + (conf?: ClientConfig) => + createClient({...clientConfig, ...(conf || {}), useDomainSharding: false}).withConfig({ + useDomainSharding: + typeof conf?.useDomainSharding === 'undefined' ? true : conf.useDomainSharding, + }), + ], + ]) + + testClient('%s: some test', (name, getClient) => { + test.skipIf(isEdge || isBrowser)( + 'can create a client that spreads request over a number of hostnames', + async () => { + const client = getClient() + + for (let i = 0; i <= 15; i++) { + const shard = i % 10 + const mockHost = `https://${defaultProjectId}.api.s${shard}.sanity.url` + const mockPath = `/v1/ping?req=${i}` + nock(mockHost).get(mockPath).delay(25).reply(200, {req: i}) + } + + const requests = [] + for (let i = 0; i <= 15; i++) { + requests.push(client.request({uri: `/ping?req=${i}`})) + } + + const responses = await Promise.all(requests) + + for (let i = 0; i <= 15; i++) { + const res = responses[i] + expect(res).toMatchObject({req: i}) + } + }, + ) + + test('listen() uses sharding', async () => { + const client = getClient() + const listenerName = 'QYdPOBgC3V0Os5QsphvTKu' + + nock('https://bf1942.api.s0.sanity.url', {encodedQueryParams: true}) + .get('/v1/data/listen/foo') + .query({query: 'true', includeResult: 'true'}) + .reply(200, `\n:\nevent: welcome\ndata: {"listenerName": "${listenerName}"}\n\n\n`, { + 'Content-Type': 'text/event-stream; charset=utf-8', + 'Transfer-Encoding': 'chunked', + }) + + return new Promise((resolve, reject) => { + const subscription = client.listen('true', {}, {events: ['welcome']}).subscribe({ + next: (msg) => { + expect(msg).toMatchObject({listenerName}) + subscription.unsubscribe() + resolve() + }, + error: (err) => { + subscription.unsubscribe() + reject(err) + }, + }) + }) + }) + + test('middleware does not shard if `useDomainSharding` is undefined', () => { + const {middleware} = getDomainSharder() + const out = middleware.processOptions({url: 'https://bf1942.api.sanity.url/v1/ping'}) + expect(out.url).toBe('https://bf1942.api.sanity.url/v1/ping') + }) + + test('middleware does not shard if `useDomainSharding` is false', () => { + const {middleware} = getDomainSharder() + const out = middleware.processOptions({ + url: 'https://bf1942.api.sanity.url/v1/ping', + useDomainSharding: false, + }) + expect(out.url).toBe('https://bf1942.api.sanity.url/v1/ping') + }) + + test('middleware rewrites hostname to be shared if `useDomainSharding` is true', () => { + const {middleware} = getDomainSharder() + const out = middleware.processOptions({ + url: 'https://bf1942.api.sanity.url/v1/ping', + useDomainSharding: true, + }) + expect(out.url).toBe('https://bf1942.api.s0.sanity.url/v1/ping') + }) + + test('middleware uses first bucket with fewest pending requests', () => { + const {middleware} = getDomainSharder([9, 6, 3, 8, 1, 2, 5, 4, 1, 7]) + const out = middleware.processOptions({ + url: 'https://bf1942.api.sanity.url/v1/ping', + useDomainSharding: true, + }) + expect(out.url).toBe('https://bf1942.api.s4.sanity.url/v1/ping') + }) + + test('middleware increases bucket request number on request', () => { + const buckets = [1, 1] + const {middleware} = getDomainSharder(buckets) + middleware.onRequest({ + options: { + url: 'https://bf1942.api.s1.sanity.url/v1/ping', + useDomainSharding: true, + }, + }) + expect(buckets).toEqual([1, 2]) + }) + + test('middleware decreases bucket request number on response', () => { + const buckets = [2, 1] + const {middleware} = getDomainSharder(buckets) + const context = { + options: { + url: 'https://bf1942.api.s0.sanity.url/v1/ping', + useDomainSharding: true, + }, + } + middleware.onResponse({} as any, context as any) + expect(buckets).toEqual([1, 1]) + }) + + test('reconfiguring with `withConfig()` maintains sharding setting', () => { + const client = getClient() + expect(client.config().useDomainSharding).toBe(true) + + const client2 = client.withConfig({apiVersion: '2024-07-01'}) + expect(client2.config().useDomainSharding).toBe(true) + }) + }) +}) From 3757e11aca370135d16d861ac20a548b90ea693a Mon Sep 17 00:00:00 2001 From: Espen Hovlandsdal Date: Tue, 2 Jul 2024 07:17:34 -0700 Subject: [PATCH 2/3] refactor(shard): use `r*` instead of `s*`, reduce to 9 buckets by default --- src/http/domainSharding.ts | 14 +++++++++----- test/domainSharding.test.ts | 14 +++++++------- 2 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/http/domainSharding.ts b/src/http/domainSharding.ts index 0b012538c..874851552 100644 --- a/src/http/domainSharding.ts +++ b/src/http/domainSharding.ts @@ -1,7 +1,8 @@ import type {Middleware, RequestOptions} from 'get-it' +const DEFAULT_NUM_SHARD_BUCKETS = 9 const UNSHARDED_URL_RE = /^https:\/\/([a-z0-9]+)\.api\.(sanity\..*)/ -const SHARDED_URL_RE = /^https:\/\/[a-z0-9]+\.api\.s(\d+)\.sanity\.(.*)/ +const SHARDED_URL_RE = /^https:\/\/[a-z0-9]+\.api\.r(\d+)\.sanity\.(.*)/ /** * Get a default sharding implementation where buckets are reused across instances. @@ -14,7 +15,7 @@ export const domainSharder = getDomainSharder() * @internal */ export function getDomainSharder(initialBuckets?: number[]) { - const buckets: number[] = initialBuckets || new Array(10).fill(0, 0) + const buckets: number[] = initialBuckets || new Array(DEFAULT_NUM_SHARD_BUCKETS).fill(0, 0) function incrementBucketForUrl(url: string) { const shard = getShardFromUrl(url) @@ -42,12 +43,16 @@ export function getDomainSharder(initialBuckets?: number[]) { 0, ) - return `https://${projectId}.api.s${bucket}.${rest}` + // We start buckets at 1, not zero - so add 1 to the bucket index + return `https://${projectId}.api.r${bucket + 1}.${rest}` } function getShardFromUrl(url: string): number | null { const [isMatch, shard] = url.match(SHARDED_URL_RE) || [] - return isMatch ? parseInt(shard, 10) : null + + // We start buckets at 1, not zero, but buckets are zero-indexed. + // Substract one from the shard number in the URL to get the correct bucket index + return isMatch ? parseInt(shard, 10) - 1 : null } const middleware = { @@ -87,7 +92,6 @@ export function getDomainSharder(initialBuckets?: number[]) { incrementBucketForUrl, decrementBucketForUrl, getShardedUrl, - getBuckets: () => buckets, } } diff --git a/test/domainSharding.test.ts b/test/domainSharding.test.ts index 59a0a1804..60a90aab3 100644 --- a/test/domainSharding.test.ts +++ b/test/domainSharding.test.ts @@ -47,8 +47,8 @@ describe('domain sharding', async () => { const client = getClient() for (let i = 0; i <= 15; i++) { - const shard = i % 10 - const mockHost = `https://${defaultProjectId}.api.s${shard}.sanity.url` + const shard = (i % 9) + 1 + const mockHost = `https://${defaultProjectId}.api.r${shard}.sanity.url` const mockPath = `/v1/ping?req=${i}` nock(mockHost).get(mockPath).delay(25).reply(200, {req: i}) } @@ -71,7 +71,7 @@ describe('domain sharding', async () => { const client = getClient() const listenerName = 'QYdPOBgC3V0Os5QsphvTKu' - nock('https://bf1942.api.s0.sanity.url', {encodedQueryParams: true}) + nock('https://bf1942.api.r1.sanity.url', {encodedQueryParams: true}) .get('/v1/data/listen/foo') .query({query: 'true', includeResult: 'true'}) .reply(200, `\n:\nevent: welcome\ndata: {"listenerName": "${listenerName}"}\n\n\n`, { @@ -115,7 +115,7 @@ describe('domain sharding', async () => { url: 'https://bf1942.api.sanity.url/v1/ping', useDomainSharding: true, }) - expect(out.url).toBe('https://bf1942.api.s0.sanity.url/v1/ping') + expect(out.url).toBe('https://bf1942.api.r1.sanity.url/v1/ping') }) test('middleware uses first bucket with fewest pending requests', () => { @@ -124,7 +124,7 @@ describe('domain sharding', async () => { url: 'https://bf1942.api.sanity.url/v1/ping', useDomainSharding: true, }) - expect(out.url).toBe('https://bf1942.api.s4.sanity.url/v1/ping') + expect(out.url).toBe('https://bf1942.api.r5.sanity.url/v1/ping') }) test('middleware increases bucket request number on request', () => { @@ -132,7 +132,7 @@ describe('domain sharding', async () => { const {middleware} = getDomainSharder(buckets) middleware.onRequest({ options: { - url: 'https://bf1942.api.s1.sanity.url/v1/ping', + url: 'https://bf1942.api.r2.sanity.url/v1/ping', useDomainSharding: true, }, }) @@ -144,7 +144,7 @@ describe('domain sharding', async () => { const {middleware} = getDomainSharder(buckets) const context = { options: { - url: 'https://bf1942.api.s0.sanity.url/v1/ping', + url: 'https://bf1942.api.r1.sanity.url/v1/ping', useDomainSharding: true, }, } From 2121beeba826c30e0702496a824d9f92bd1c6741 Mon Sep 17 00:00:00 2001 From: Espen Hovlandsdal Date: Tue, 2 Jul 2024 14:23:25 -0700 Subject: [PATCH 3/3] test(shard): skip listen test in browsers/edge --- test/domainSharding.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/domainSharding.test.ts b/test/domainSharding.test.ts index 60a90aab3..d9cfe4b45 100644 --- a/test/domainSharding.test.ts +++ b/test/domainSharding.test.ts @@ -67,7 +67,7 @@ describe('domain sharding', async () => { }, ) - test('listen() uses sharding', async () => { + test.skipIf(isEdge || isBrowser)('listen() uses sharding', async () => { const client = getClient() const listenerName = 'QYdPOBgC3V0Os5QsphvTKu'