diff --git a/packages/typescript/src/api/async/client.ts b/packages/typescript/src/api/async/client.ts index 5ed5397cf55cd..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); } @@ -214,10 +220,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 +276,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..1534850abb11c 100644 --- a/packages/typescript/test/sync/api-generators.test.ts +++ b/packages/typescript/test/sync/api-generators.test.ts @@ -481,6 +481,47 @@ 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("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/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..90d0128f7db09 100644 --- a/tsc/internal/api/session.go +++ b/tsc/internal/api/session.go @@ -371,7 +371,9 @@ type Session struct { projectSession *project.Session // This is set to true when using MessagePackProtocol. - useBinaryResponses bool + useBinaryResponses bool + batchResponsePages sync.Map + nextBatchResponsePageID atomic.Uint64 // snapshots maps snapshot handles to their data. Each snapshot has its own // symbol/type registries. @@ -411,6 +413,10 @@ type Session struct { cpuProfiler pprof.CPUProfiler } +type batchResponsePage struct { + encodedResponses []json.Value +} + // Ensure Session implements Handler var _ ipc.Handler = (*Session)(nil) @@ -420,6 +426,10 @@ 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) @@ -890,15 +900,90 @@ 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 != "" { + 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) + } + 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.batchResponsePages.Store(continuationToken, batchResponsePage{ + encodedResponses: slices.Clone(page.encodedResponses[pageLength:]), + }) + 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 +3737,7 @@ func computeSnapshotChanges(prev *project.Snapshot, next *project.Snapshot) *Sna // regardless of their ref counts. func (s *Session) Close() { s.releaseOpenRefs() + 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 1a57c51d52405..e8866c062f40a 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)) + + value, ok := session.batchResponsePages.Load(response.ContinuationToken) + assert.Assert(t, ok) + pending := value.(batchResponsePage) + 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")) +}