From 708f8666a42c2257cfb5045f584e130aeb32019e Mon Sep 17 00:00:00 2001 From: Wesley Wigham Date: Thu, 27 Aug 2026 12:37:03 -0700 Subject: [PATCH 1/2] Add pagination of batch requests --- packages/typescript/src/api/async/client.ts | 21 +++- packages/typescript/src/api/options.ts | 4 + .../typescript/src/api/proto.generated.ts | 3 + packages/typescript/src/api/proto.ts | 3 +- packages/typescript/src/api/sync/client.ts | 24 +++- packages/typescript/test/async/api.test.ts | 22 ++++ .../test/sync/api-generators.test.ts | 18 +++ tsc/internal/api/proto.go | 47 ++++++- tsc/internal/api/session.go | 107 +++++++++++++++- tsc/internal/api/session_batch_test.go | 117 ++++++++++++++++++ 10 files changed, 355 insertions(+), 11 deletions(-) diff --git a/packages/typescript/src/api/async/client.ts b/packages/typescript/src/api/async/client.ts index 5ed5397cf55cd..9c5ce6624cf9b 100644 --- a/packages/typescript/src/api/async/client.ts +++ b/packages/typescript/src/api/async/client.ts @@ -214,10 +214,27 @@ export class Client { const requestType = new RequestType("batchRequests"); const params: BatchRequestsParams = { requests: requests.map(request => ({ method: request.method, params: request.params })) }; + if (this.options.maxResponseBytesPerPage !== undefined) { + params.maxResponseBytesPerPage = this.options.maxResponseBytesPerPage; + } const response = await this.sendRequestWithTiming(requestType, params); + let responses = response.responses; + let continuationToken = response.continuationToken; + while (continuationToken) { + const pageParams: BatchRequestsParams = { + requests: [], + continuationToken, + }; + if (this.options.maxResponseBytesPerPage !== undefined) { + pageParams.maxResponseBytesPerPage = this.options.maxResponseBytesPerPage; + } + const page = await this.sendRequestWithTiming(requestType, pageParams); + responses = responses.concat(page.responses); + continuationToken = page.continuationToken; + } for (let i = 0; i < requests.length; i++) { const { resolve, reject } = requests[i]; - const item = response.responses[i]; + const item = responses[i]; if (item.error !== undefined) { reject(new Error(item.error)); } @@ -253,7 +270,7 @@ export class Client { }; } - async apiRequest(method: K, params: APIMethodInfo[K]["params"]): Promise { + async apiRequest(method: K, params: APIMethodInfo[K]["params"]): Promise { if (!this.connected) { await this.connect(); } diff --git a/packages/typescript/src/api/options.ts b/packages/typescript/src/api/options.ts index 759523393d0f2..07e73941511b0 100644 --- a/packages/typescript/src/api/options.ts +++ b/packages/typescript/src/api/options.ts @@ -8,6 +8,8 @@ import type { FileSystem } from "./fs.ts"; export interface ClientSocketOptions { /** Path to the Unix domain socket or Windows named pipe for API communication */ pipe: string; + /** Maximum encoded byte size of each batch response page. Defaults to 300 million bytes. */ + maxResponseBytesPerPage?: number; } export interface ClientSpawnOptions { @@ -19,6 +21,8 @@ export interface ClientSpawnOptions { fs?: FileSystem; /** Allow trusted projects to execute configured external content mapper processes. */ runExternalCode?: boolean; + /** Maximum encoded byte size of each batch response page. Defaults to 300 million bytes. */ + maxResponseBytesPerPage?: number; /** * When true, collect timing information for each request. The client * measures round-trip latency and bytes sent/received, and the server diff --git a/packages/typescript/src/api/proto.generated.ts b/packages/typescript/src/api/proto.generated.ts index 6528b59b0f3f9..a6e2edf37c1c8 100644 --- a/packages/typescript/src/api/proto.generated.ts +++ b/packages/typescript/src/api/proto.generated.ts @@ -163,10 +163,13 @@ export interface ReleaseParams { export interface BatchRequestsParams { requests: readonly BatchRequest[] | null; + continuationToken?: string; + maxResponseBytesPerPage?: number; } export interface BatchRequestsResponse { responses: BatchResponse[]; + continuationToken?: string; } /** InitializeResponse is returned by the initialize method. */ diff --git a/packages/typescript/src/api/proto.ts b/packages/typescript/src/api/proto.ts index 6c29b4c3a4a13..f3e7cc00d5e35 100644 --- a/packages/typescript/src/api/proto.ts +++ b/packages/typescript/src/api/proto.ts @@ -25,7 +25,8 @@ export type TypePropertyMethod = Exclude, Intr export type TypesPropertyMethod = APIMethodsReturning; export type IntrinsicTypeMethod = "getAnyType" | "getBigIntType" | "getBooleanType" | "getESSymbolType" | "getNeverType" | "getNonPrimitiveType" | "getNullType" | "getNumberType" | "getStringType" | "getUndefinedType" | "getUnknownType" | "getVoidType"; -export type APIRequest = { [K in keyof APIMethodInfo]: { method: K; params: APIMethodInfo[K]["params"]; }; }[keyof APIMethodInfo]; +type BatchableAPIMethod = Exclude; +export type APIRequest = { [K in BatchableAPIMethod]: { method: K; params: APIMethodInfo[K]["params"]; }; }[BatchableAPIMethod]; export type APIResponse = Request extends APIRequest ? & { method: Request["method"]; diff --git a/packages/typescript/src/api/sync/client.ts b/packages/typescript/src/api/sync/client.ts index 027e30f8a57d0..ca22e5009a466 100644 --- a/packages/typescript/src/api/sync/client.ts +++ b/packages/typescript/src/api/sync/client.ts @@ -10,6 +10,7 @@ import { import type { APIMethodInfo, APIRequest, + BatchRequestsParams, BatchRequestsResponse, SourceFileResponseMethod, } from "../proto.ts"; @@ -28,6 +29,7 @@ export class Client { private channel: SyncRpcChannel; private encoder = new TextEncoder(); private timing: TimingCollector | undefined; + private maxResponseBytesPerPage: number | undefined; constructor(options: ClientOptions) { if (!isSpawnOptions(options)) { @@ -35,6 +37,7 @@ export class Client { } const args = getAPIProcessArgs(options, false); + this.maxResponseBytesPerPage = options.maxResponseBytesPerPage; // Enable virtual FS callbacks for each provided FS function const enabledCallbacks: (typeof fsCallbackNames[number])[] = []; @@ -99,7 +102,26 @@ export class Client { } batchRequests(requests: readonly APIRequest[]): BatchRequestsResponse { - return this.apiRequest("batchRequests", { requests }); + const params: BatchRequestsParams = { requests }; + if (this.maxResponseBytesPerPage !== undefined) { + params.maxResponseBytesPerPage = this.maxResponseBytesPerPage; + } + const response = this.apiRequest("batchRequests", params); + let responses = response.responses; + let continuationToken = response.continuationToken; + while (continuationToken) { + const pageParams: BatchRequestsParams = { + requests: [], + continuationToken, + }; + if (this.maxResponseBytesPerPage !== undefined) { + pageParams.maxResponseBytesPerPage = this.maxResponseBytesPerPage; + } + const page = this.apiRequest("batchRequests", pageParams); + responses = responses.concat(page.responses); + continuationToken = page.continuationToken; + } + return { responses }; } apiRequestBinary(method: K, params?: APIMethodInfo[K]["params"]): Uint8Array | undefined { diff --git a/packages/typescript/test/async/api.test.ts b/packages/typescript/test/async/api.test.ts index b40091c90d6c9..6c19ab32d1795 100644 --- a/packages/typescript/test/async/api.test.ts +++ b/packages/typescript/test/async/api.test.ts @@ -421,6 +421,28 @@ describe("API - automatic batching", () => { }); describe("API - batchContext", () => { + test("transparently paginates batch responses", async () => { + const api = spawnAPI({ ...defaultFiles }, { maxResponseBytesPerPage: 1 }); + try { + const requests = await (async () => { + using _ = api.batchContext(); + return [ + api.parseCommandLine(["--strict"]), + api.readConfigFile("/tsconfig.json"), + api.parseCommandLine(["--noImplicitAny"]), + ] as const; + })(); + + const [strict, config, noImplicitAny] = await Promise.all(requests); + assert.equal(strict.options.strict, true); + assert.deepEqual(config.config, {}); + assert.equal(noImplicitAny.options.noImplicitAny, true); + } + finally { + await api.close(); + } + }); + test("holds requests until disposal", async () => { const api = spawnAPI(); try { diff --git a/packages/typescript/test/sync/api-generators.test.ts b/packages/typescript/test/sync/api-generators.test.ts index 9016e6c75e7de..f5c32c99646be 100644 --- a/packages/typescript/test/sync/api-generators.test.ts +++ b/packages/typescript/test/sync/api-generators.test.ts @@ -481,6 +481,24 @@ describe("API - generator batching", () => { } }); + test("transparently paginates batch responses", () => { + const api = spawnAPI(undefined, { maxResponseBytesPerPage: 1 }); + try { + const [strict, config, noImplicitAny] = api.batch( + api.parseCommandLine.gen(["--strict"]), + api.readConfigFile.gen("/tsconfig.json"), + api.parseCommandLine.gen(["--noImplicitAny"]), + ); + + assert.equal(strict.options.strict, true); + assert.deepEqual(config.config, {}); + assert.equal(noImplicitAny.options.noImplicitAny, true); + } + finally { + api.close(); + } + }); + test("all deduplicates only initialize requests within a batch round", () => { const api = spawnAPI(); const requestBatches: string[][] = []; diff --git a/tsc/internal/api/proto.go b/tsc/internal/api/proto.go index 84b3673425fce..7321c1877481b 100644 --- a/tsc/internal/api/proto.go +++ b/tsc/internal/api/proto.go @@ -621,7 +621,9 @@ type TranspileOutputResponse struct { } type BatchRequestsParams struct { - Requests []BatchRequest `json:"requests"` + Requests []BatchRequest `json:"requests"` + ContinuationToken string `json:"continuationToken,omitempty"` + MaxResponseBytesPerPage int `json:"maxResponseBytesPerPage,omitempty"` } type BatchRequest struct { @@ -630,7 +632,48 @@ type BatchRequest struct { } type BatchRequestsResponse struct { - Responses []BatchResponse `json:"responses" nonnil:"true"` + Responses []BatchResponse `json:"responses" nonnil:"true"` + ContinuationToken string `json:"continuationToken,omitempty"` + encodedResponses []json.Value +} + +var _ json.MarshalerTo = (*BatchRequestsResponse)(nil) + +func (r *BatchRequestsResponse) MarshalJSONTo(enc *json.Encoder) error { + if err := enc.WriteToken(json.BeginObject); err != nil { + return err + } + if err := enc.WriteValue(json.Value(`"responses"`)); err != nil { + return err + } + if err := enc.WriteToken(json.BeginArray); err != nil { + return err + } + if r.encodedResponses != nil { + for _, response := range r.encodedResponses { + if err := enc.WriteValue(response); err != nil { + return err + } + } + } else { + for i := range r.Responses { + if err := json.MarshalEncode(enc, &r.Responses[i]); err != nil { + return err + } + } + } + if err := enc.WriteToken(json.EndArray); err != nil { + return err + } + if r.ContinuationToken != "" { + if err := enc.WriteValue(json.Value(`"continuationToken"`)); err != nil { + return err + } + if err := json.MarshalEncode(enc, r.ContinuationToken); err != nil { + return err + } + } + return enc.WriteToken(json.EndObject) } type BatchResponse struct { diff --git a/tsc/internal/api/session.go b/tsc/internal/api/session.go index 03c28a6d665e1..1221473295cf3 100644 --- a/tsc/internal/api/session.go +++ b/tsc/internal/api/session.go @@ -371,7 +371,10 @@ type Session struct { projectSession *project.Session // This is set to true when using MessagePackProtocol. - useBinaryResponses bool + useBinaryResponses bool + batchResponsePages map[string]batchResponsePage + batchResponsePagesMu sync.Mutex + nextBatchResponsePageID atomic.Uint64 // snapshots maps snapshot handles to their data. Each snapshot has its own // symbol/type registries. @@ -411,6 +414,10 @@ type Session struct { cpuProfiler pprof.CPUProfiler } +type batchResponsePage struct { + encodedResponses []json.Value +} + // Ensure Session implements Handler var _ ipc.Handler = (*Session)(nil) @@ -420,13 +427,18 @@ type SessionOptions struct { UseBinaryResponses bool } +// DefaultMaxResponseBytesPerPage leaves room for base64 expansion beneath V8's +// maximum string length while rounding down to an even decimal value. +const DefaultMaxResponseBytesPerPage = 300_000_000 + // NewSession creates a new API session with the given project session. func NewSession(projectSession *project.Session, options *SessionOptions) *Session { id := sessionIDCounter.Add(1) s := &Session{ - id: formatSessionID(id), - projectSession: projectSession, - snapshots: make(map[SnapshotID]*snapshotData), + id: formatSessionID(id), + projectSession: projectSession, + batchResponsePages: make(map[string]batchResponsePage), + snapshots: make(map[SnapshotID]*snapshotData), } if options != nil { s.useBinaryResponses = options.UseBinaryResponses @@ -890,15 +902,97 @@ func (s *Session) HandleRequest(ctx context.Context, method string, params json. } func (s *Session) handleBatchRequests(ctx context.Context, params *BatchRequestsParams) (*BatchRequestsResponse, error) { + if params.ContinuationToken != "" { + s.batchResponsePagesMu.Lock() + page, ok := s.batchResponsePages[params.ContinuationToken] + delete(s.batchResponsePages, params.ContinuationToken) + s.batchResponsePagesMu.Unlock() + if !ok { + return nil, fmt.Errorf("%w: invalid batch continuation token", ErrClientError) + } + return s.paginateBatchResponses(page, nil, params.MaxResponseBytesPerPage) + } + responses := make([]BatchResponse, len(params.Requests)) for i, request := range params.Requests { responses[i] = s.handleBatchRequest(ctx, request) } - return &BatchRequestsResponse{Responses: responses}, nil + if s == nil { + return &BatchRequestsResponse{Responses: responses}, nil + } + page, err := newBatchResponsePage(responses) + if err != nil { + return nil, err + } + return s.paginateBatchResponses(page, responses, params.MaxResponseBytesPerPage) +} + +func newBatchResponsePage(responses []BatchResponse) (batchResponsePage, error) { + encodedResponses := make([]json.Value, len(responses)) + for i := range responses { + encoded, err := json.Marshal(&responses[i]) + if err != nil { + return batchResponsePage{}, err + } + encodedResponses[i] = encoded + } + return batchResponsePage{encodedResponses: encodedResponses}, nil +} + +func (s *Session) paginateBatchResponses(page batchResponsePage, responses []BatchResponse, maxResponseBytesPerPage int) (*BatchRequestsResponse, error) { + if maxResponseBytesPerPage <= 0 { + maxResponseBytesPerPage = DefaultMaxResponseBytesPerPage + } + encodedLength := len(`{"responses":[]}`) + pageLength := 0 + for _, encoded := range page.encodedResponses { + additionalLength := len(encoded) + if pageLength > 0 { + additionalLength++ + } + if pageLength > 0 && encodedLength+additionalLength > maxResponseBytesPerPage { + break + } + encodedLength += additionalLength + pageLength++ + } + + if pageLength == len(page.encodedResponses) { + return &BatchRequestsResponse{ + Responses: responses, + encodedResponses: page.encodedResponses, + }, nil + } + continuationToken := fmt.Sprintf("%s-%d", s.id, s.nextBatchResponsePageID.Add(1)) + continuationLength := len(`,"continuationToken":""`) + len(continuationToken) + for pageLength > 1 && encodedLength+continuationLength > maxResponseBytesPerPage { + encodedLength -= len(page.encodedResponses[pageLength-1]) + 1 + pageLength-- + } + response := &BatchRequestsResponse{ + ContinuationToken: continuationToken, + encodedResponses: slices.Clone(page.encodedResponses[:pageLength]), + } + if responses != nil { + response.Responses = slices.Clone(responses[:pageLength]) + } + s.batchResponsePagesMu.Lock() + if s.batchResponsePages == nil { + s.batchResponsePages = make(map[string]batchResponsePage) + } + s.batchResponsePages[continuationToken] = batchResponsePage{ + encodedResponses: slices.Clone(page.encodedResponses[pageLength:]), + } + s.batchResponsePagesMu.Unlock() + return response, nil } func (s *Session) handleBatchRequest(ctx context.Context, request BatchRequest) (response BatchResponse) { response.Method = request.Method + if request.Method == MethodBatchRequests { + response.Error = fmt.Sprintf("%s: batchRequests cannot be nested", ErrInvalidRequest) + return response + } defer func() { if recovered := recover(); recovered != nil { response.Result = nil @@ -3652,6 +3746,9 @@ func computeSnapshotChanges(prev *project.Snapshot, next *project.Snapshot) *Sna // regardless of their ref counts. func (s *Session) Close() { s.releaseOpenRefs() + s.batchResponsePagesMu.Lock() + clear(s.batchResponsePages) + s.batchResponsePagesMu.Unlock() s.snapshotsMu.Lock() defer s.snapshotsMu.Unlock() diff --git a/tsc/internal/api/session_batch_test.go b/tsc/internal/api/session_batch_test.go index 1a57c51d52405..8ba51ac6b22ea 100644 --- a/tsc/internal/api/session_batch_test.go +++ b/tsc/internal/api/session_batch_test.go @@ -59,3 +59,120 @@ func TestBatchResponseEncodesEmptyResult(t *testing.T) { assert.NilError(t, err) assert.Equal(t, string(encoded), `{"method":"getSignaturesOfType","result":[]}`) } + +func TestHandleBatchRequestsPaginatesResponses(t *testing.T) { + t.Parallel() + + const maxResponseBytesPerPage = 150 + session := NewSession(nil, nil) + requests := make([]BatchRequest, 10) + for i := range requests { + requests[i] = BatchRequest{Method: "ping", Params: json.Value{}} + } + + response, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{ + Requests: requests, + MaxResponseBytesPerPage: maxResponseBytesPerPage, + }) + assert.NilError(t, err) + var responses []BatchResponse + for { + encoded, err := json.Marshal(response) + assert.NilError(t, err) + assert.Assert(t, len(encoded) <= maxResponseBytesPerPage) + var wireResponse BatchRequestsResponse + assert.NilError(t, json.Unmarshal(encoded, &wireResponse)) + responses = append(responses, wireResponse.Responses...) + if wireResponse.ContinuationToken == "" { + break + } + response, err = session.handleBatchRequests(context.Background(), &BatchRequestsParams{ + ContinuationToken: wireResponse.ContinuationToken, + MaxResponseBytesPerPage: maxResponseBytesPerPage, + }) + assert.NilError(t, err) + } + + assert.Equal(t, len(responses), len(requests)) + for _, response := range responses { + assert.Equal(t, response.Method, Method("ping")) + assert.Equal(t, response.Result, "pong") + } +} + +func TestHandleBatchRequestsAllowsOversizedSingleResponse(t *testing.T) { + t.Parallel() + + session := NewSession(nil, nil) + response, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{ + Requests: []BatchRequest{{Method: "ping", Params: json.Value{}}}, + MaxResponseBytesPerPage: 1, + }) + assert.NilError(t, err) + assert.Equal(t, len(response.Responses), 1) + assert.Equal(t, response.ContinuationToken, "") +} + +func TestHandleBatchRequestsPageLimitIsRequestScoped(t *testing.T) { + t.Parallel() + + session := NewSession(nil, nil) + requests := []BatchRequest{ + {Method: "ping", Params: json.Value{}}, + {Method: "ping", Params: json.Value{}}, + } + + limited, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{ + Requests: requests, + MaxResponseBytesPerPage: 1, + }) + assert.NilError(t, err) + assert.Equal(t, len(limited.Responses), 1) + assert.Assert(t, limited.ContinuationToken != "") + + unlimited, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{Requests: requests}) + assert.NilError(t, err) + assert.Equal(t, len(unlimited.Responses), len(requests)) + assert.Equal(t, unlimited.ContinuationToken, "") +} + +func TestHandleBatchRequestsDoesNotRetainSentPageBackingArray(t *testing.T) { + t.Parallel() + + session := NewSession(nil, nil) + response, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{ + Requests: []BatchRequest{ + {Method: "ping", Params: json.Value{}}, + {Method: "ping", Params: json.Value{}}, + {Method: "ping", Params: json.Value{}}, + }, + MaxResponseBytesPerPage: 1, + }) + assert.NilError(t, err) + assert.Equal(t, cap(response.encodedResponses), len(response.encodedResponses)) + + session.batchResponsePagesMu.Lock() + pending := session.batchResponsePages[response.ContinuationToken] + session.batchResponsePagesMu.Unlock() + assert.Equal(t, cap(pending.encodedResponses), len(pending.encodedResponses)) +} + +func TestHandleBatchRequestsRejectsInvalidContinuationToken(t *testing.T) { + t.Parallel() + + session := NewSession(nil, nil) + _, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{ContinuationToken: "invalid"}) + assert.ErrorContains(t, err, "invalid batch continuation token") +} + +func TestHandleBatchRequestsRejectsNestedBatch(t *testing.T) { + t.Parallel() + + session := NewSession(nil, nil) + response, err := session.handleBatchRequests(context.Background(), &BatchRequestsParams{ + Requests: []BatchRequest{{Method: MethodBatchRequests, Params: json.Value(`{"requests":[]}`)}}, + }) + assert.NilError(t, err) + assert.Equal(t, len(response.Responses), 1) + assert.Assert(t, strings.Contains(response.Responses[0].Error, "batchRequests cannot be nested")) +} From 7100d4333b458314bd97ec1b7b18dee5236600b3 Mon Sep 17 00:00:00 2001 From: Wesley Wigham Date: Thu, 27 Aug 2026 15:32:09 -0700 Subject: [PATCH 2/2] Add large-response stress test of the default batch size, fix hang in async client it exposed --- packages/typescript/src/api/async/client.ts | 10 +++++-- .../test/sync/api-generators.test.ts | 23 +++++++++++++++ tsc/internal/api/session.go | 29 ++++++------------- tsc/internal/api/session_batch_test.go | 6 ++-- 4 files changed, 43 insertions(+), 25 deletions(-) diff --git a/packages/typescript/src/api/async/client.ts b/packages/typescript/src/api/async/client.ts index 9c5ce6624cf9b..e4ba089f4e861 100644 --- a/packages/typescript/src/api/async/client.ts +++ b/packages/typescript/src/api/async/client.ts @@ -49,6 +49,7 @@ export class Client { private connection: MessageConnection | undefined; private options: ClientOptions; private connected = false; + private connecting: Promise | undefined; private timing: TimingCollector | undefined; private batchedRequests: { method: APIRequest["method"]; params: APIRequest["params"]; resolve: (value: unknown) => void; reject: (reason?: any) => void; }[] = []; private nextBatch: NodeJS.Immediate | "manual" | undefined; @@ -60,9 +61,14 @@ export class Client { } } - async connect(): Promise { - if (this.connected) return; + connect(): Promise { + if (this.connected) return Promise.resolve(); + return this.connecting ??= this.connectWorker().finally(() => { + this.connecting = undefined; + }); + } + private async connectWorker(): Promise { if (isSpawnOptions(this.options)) { await this.connectViaSpawn(this.options); } diff --git a/packages/typescript/test/sync/api-generators.test.ts b/packages/typescript/test/sync/api-generators.test.ts index f5c32c99646be..1534850abb11c 100644 --- a/packages/typescript/test/sync/api-generators.test.ts +++ b/packages/typescript/test/sync/api-generators.test.ts @@ -499,6 +499,29 @@ describe("API - generator batching", () => { } }); + test("transparently paginates responses at the default batch size limit", () => { + const largeConfigValue = "x".repeat(5_000_000); + const requestCount = 64; + const api = spawnAPI({ "/large.json": JSON.stringify({ largeConfigValue }) }, { collectTiming: true }); + try { + api.parseCommandLine([]); + api.resetTimingInfo(); + + const configs = api.batch(...Array.from({ length: requestCount }, () => api.readConfigFile.gen("/large.json"))); + assert.equal(configs.length, requestCount); + for (const config of configs) { + assert.deepEqual(config.config, { largeConfigValue }); + } + + const timing = api.getTimingInfo(); + assert.equal(timing.totals.requestCount, 2); + assert.deepEqual(timing.recentRequests.map(request => request.method), ["batchRequests", "batchRequests"]); + } + finally { + api.close(); + } + }); + test("all deduplicates only initialize requests within a batch round", () => { const api = spawnAPI(); const requestBatches: string[][] = []; diff --git a/tsc/internal/api/session.go b/tsc/internal/api/session.go index 1221473295cf3..90d0128f7db09 100644 --- a/tsc/internal/api/session.go +++ b/tsc/internal/api/session.go @@ -372,8 +372,7 @@ type Session struct { // This is set to true when using MessagePackProtocol. useBinaryResponses bool - batchResponsePages map[string]batchResponsePage - batchResponsePagesMu sync.Mutex + batchResponsePages sync.Map nextBatchResponsePageID atomic.Uint64 // snapshots maps snapshot handles to their data. Each snapshot has its own @@ -435,10 +434,9 @@ const DefaultMaxResponseBytesPerPage = 300_000_000 func NewSession(projectSession *project.Session, options *SessionOptions) *Session { id := sessionIDCounter.Add(1) s := &Session{ - id: formatSessionID(id), - projectSession: projectSession, - batchResponsePages: make(map[string]batchResponsePage), - snapshots: make(map[SnapshotID]*snapshotData), + id: formatSessionID(id), + projectSession: projectSession, + snapshots: make(map[SnapshotID]*snapshotData), } if options != nil { s.useBinaryResponses = options.UseBinaryResponses @@ -903,13 +901,11 @@ func (s *Session) HandleRequest(ctx context.Context, method string, params json. func (s *Session) handleBatchRequests(ctx context.Context, params *BatchRequestsParams) (*BatchRequestsResponse, error) { if params.ContinuationToken != "" { - s.batchResponsePagesMu.Lock() - page, ok := s.batchResponsePages[params.ContinuationToken] - delete(s.batchResponsePages, params.ContinuationToken) - s.batchResponsePagesMu.Unlock() + value, ok := s.batchResponsePages.LoadAndDelete(params.ContinuationToken) if !ok { return nil, fmt.Errorf("%w: invalid batch continuation token", ErrClientError) } + page := value.(batchResponsePage) return s.paginateBatchResponses(page, nil, params.MaxResponseBytesPerPage) } @@ -976,14 +972,9 @@ func (s *Session) paginateBatchResponses(page batchResponsePage, responses []Bat if responses != nil { response.Responses = slices.Clone(responses[:pageLength]) } - s.batchResponsePagesMu.Lock() - if s.batchResponsePages == nil { - s.batchResponsePages = make(map[string]batchResponsePage) - } - s.batchResponsePages[continuationToken] = batchResponsePage{ + s.batchResponsePages.Store(continuationToken, batchResponsePage{ encodedResponses: slices.Clone(page.encodedResponses[pageLength:]), - } - s.batchResponsePagesMu.Unlock() + }) return response, nil } @@ -3746,9 +3737,7 @@ func computeSnapshotChanges(prev *project.Snapshot, next *project.Snapshot) *Sna // regardless of their ref counts. func (s *Session) Close() { s.releaseOpenRefs() - s.batchResponsePagesMu.Lock() - clear(s.batchResponsePages) - s.batchResponsePagesMu.Unlock() + s.batchResponsePages.Clear() s.snapshotsMu.Lock() defer s.snapshotsMu.Unlock() diff --git a/tsc/internal/api/session_batch_test.go b/tsc/internal/api/session_batch_test.go index 8ba51ac6b22ea..e8866c062f40a 100644 --- a/tsc/internal/api/session_batch_test.go +++ b/tsc/internal/api/session_batch_test.go @@ -151,9 +151,9 @@ func TestHandleBatchRequestsDoesNotRetainSentPageBackingArray(t *testing.T) { assert.NilError(t, err) assert.Equal(t, cap(response.encodedResponses), len(response.encodedResponses)) - session.batchResponsePagesMu.Lock() - pending := session.batchResponsePages[response.ContinuationToken] - session.batchResponsePagesMu.Unlock() + value, ok := session.batchResponsePages.Load(response.ContinuationToken) + assert.Assert(t, ok) + pending := value.(batchResponsePage) assert.Equal(t, cap(pending.encodedResponses), len(pending.encodedResponses)) }