Skip to content

Commit 34c9cae

Browse files
author
Benjamin Pasero
committed
files - avoid stream transform when peeking into stream
1 parent 4002404 commit 34c9cae

5 files changed

Lines changed: 187 additions & 80 deletions

File tree

‎src/vs/base/common/buffer.ts‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,8 @@ export interface VSBufferReadableStream extends streams.ReadableStream<VSBuffer>
194194

195195
export interface VSBufferWriteableStream extends streams.WriteableStream<VSBuffer> { }
196196

197+
export interface VSBufferReadableBufferedStream extends streams.ReadableBufferedStream<VSBuffer> { }
198+
197199
export function readableToBuffer(readable: VSBufferReadable): VSBuffer {
198200
return streams.consumeReadable<VSBuffer>(readable, chunks => VSBuffer.concat(chunks));
199201
}
@@ -206,6 +208,21 @@ export function streamToBuffer(stream: streams.ReadableStream<VSBuffer>): Promis
206208
return streams.consumeStream<VSBuffer>(stream, chunks => VSBuffer.concat(chunks));
207209
}
208210

211+
export async function bufferedStreamToBuffer(bufferedStream: streams.ReadableBufferedStream<VSBuffer>): Promise<VSBuffer> {
212+
if (bufferedStream.ended) {
213+
return VSBuffer.concat(bufferedStream.buffer);
214+
}
215+
216+
return VSBuffer.concat([
217+
218+
// Include already read chunks...
219+
...bufferedStream.buffer,
220+
221+
// ...and all additional chunks
222+
await streamToBuffer(bufferedStream.stream)
223+
]);
224+
}
225+
209226
export function bufferToStream(buffer: VSBuffer): streams.ReadableStream<VSBuffer> {
210227
return streams.toStream<VSBuffer>(buffer, chunks => VSBuffer.concat(chunks));
211228
}

‎src/vs/base/common/stream.ts‎

Lines changed: 66 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6+
import { DisposableStore, toDisposable } from 'vs/base/common/lifecycle';
7+
68
/**
79
* The payload that flows in readable stream events.
810
*/
@@ -106,12 +108,43 @@ export interface WriteableStream<T> extends ReadableStream<T> {
106108
end(result?: T | Error): void;
107109
}
108110

111+
/**
112+
* A stream that has a buffer already read. Returns the original stream
113+
* that was read as well as the chunks that got read.
114+
*
115+
* The `ended` flag indicates if the stream has been fully consumed.
116+
*/
117+
export interface ReadableBufferedStream<T> {
118+
119+
/**
120+
* The original stream that is being read.
121+
*/
122+
stream: ReadableStream<T>;
123+
124+
/**
125+
* An array of chunks already read from this stream.
126+
*/
127+
buffer: T[];
128+
129+
/**
130+
* Signals if the stream has ended or not. If not, consumers
131+
* should continue to read from the stream until consumed.
132+
*/
133+
ended: boolean;
134+
}
135+
109136
export function isReadableStream<T>(obj: unknown): obj is ReadableStream<T> {
110137
const candidate = obj as ReadableStream<T>;
111138

112139
return candidate && [candidate.on, candidate.pause, candidate.resume, candidate.destroy].every(fn => typeof fn === 'function');
113140
}
114141

142+
export function isReadableBufferedStream<T>(obj: unknown): obj is ReadableBufferedStream<T> {
143+
const candidate = obj as ReadableBufferedStream<T>;
144+
145+
return candidate && isReadableStream(candidate.stream) && Array.isArray(candidate.buffer) && typeof candidate.ended === 'boolean';
146+
}
147+
115148
export interface IReducer<T> {
116149
(data: T[]): T;
117150
}
@@ -395,7 +428,7 @@ export function consumeReadable<T>(readable: Readable<T>, reducer: IReducer<T>):
395428
* reached, will return a readable instead to ensure all data can still
396429
* be read.
397430
*/
398-
export function consumeReadableWithLimit<T>(readable: Readable<T>, reducer: IReducer<T>, maxChunks: number): T | Readable<T> {
431+
export function peekReadable<T>(readable: Readable<T>, reducer: IReducer<T>, maxChunks: number): T | Readable<T> {
399432
const chunks: T[] = [];
400433

401434
let chunk: T | null | undefined = undefined;
@@ -452,58 +485,50 @@ export function consumeStream<T>(stream: ReadableStream<T>, reducer: IReducer<T>
452485
}
453486

454487
/**
455-
* Helper to read a T stream up to a maximum of chunks. If the limit is
456-
* reached, will return a stream instead to ensure all data can still
457-
* be read.
488+
* Helper to peek up to `maxChunks` into a stream. The return type signals if
489+
* the stream has ended or not. If not, caller needs to add a `data` listener
490+
* to continue reading.
458491
*/
459-
export function consumeStreamWithLimit<T>(stream: ReadableStream<T>, reducer: IReducer<T>, maxChunks: number): Promise<T | ReadableStream<T>> {
492+
export function peekStream<T>(stream: ReadableStream<T>, maxChunks: number): Promise<ReadableBufferedStream<T>> {
460493
return new Promise((resolve, reject) => {
461-
const chunks: T[] = [];
494+
const streamListeners = new DisposableStore();
462495

463-
let wrapperStream: WriteableStream<T> | undefined = undefined;
496+
// Data Listener
497+
const buffer: T[] = [];
498+
const dataListener = (chunk: T) => {
464499

465-
stream.on('data', data => {
500+
// Add to buffer
501+
buffer.push(chunk);
466502

467-
// If we reach maxChunks, we start to return a stream
468-
// and make sure that any data we have already read
469-
// is in it as well
470-
if (!wrapperStream && chunks.length === maxChunks) {
471-
wrapperStream = newWriteableStream(reducer);
472-
while (chunks.length) {
473-
wrapperStream.write(chunks.shift()!);
474-
}
503+
// We reached maxChunks and thus need to return
504+
if (buffer.length > maxChunks) {
475505

476-
wrapperStream.write(data);
506+
// Dispose any listeners and ensure to pause the
507+
// stream so that it can be consumed again by caller
508+
streamListeners.dispose();
509+
stream.pause();
477510

478-
return resolve(wrapperStream);
511+
return resolve({ stream, buffer, ended: false });
479512
}
513+
};
480514

481-
if (wrapperStream) {
482-
wrapperStream.write(data);
483-
} else {
484-
chunks.push(data);
485-
}
486-
});
515+
streamListeners.add(toDisposable(() => stream.removeListener('data', dataListener)));
516+
stream.on('data', dataListener);
487517

488-
stream.on('error', error => {
489-
if (wrapperStream) {
490-
wrapperStream.error(error);
491-
} else {
492-
return reject(error);
493-
}
494-
});
518+
// Error Listener
519+
const errorListener = (error: Error) => {
520+
return reject(error);
521+
};
495522

496-
stream.on('end', () => {
497-
if (wrapperStream) {
498-
while (chunks.length) {
499-
wrapperStream.write(chunks.shift()!);
500-
}
523+
streamListeners.add(toDisposable(() => stream.removeListener('error', errorListener)));
524+
stream.on('error', errorListener);
501525

502-
wrapperStream.end();
503-
} else {
504-
return resolve(reducer(chunks));
505-
}
506-
});
526+
const endListener = () => {
527+
return resolve({ stream, buffer, ended: true });
528+
};
529+
530+
streamListeners.add(toDisposable(() => stream.removeListener('end', endListener)));
531+
stream.on('end', endListener);
507532
});
508533
}
509534

‎src/vs/base/test/common/stream.test.ts‎

Lines changed: 38 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import * as assert from 'assert';
7-
import { isReadableStream, newWriteableStream, Readable, consumeReadable, consumeReadableWithLimit, consumeStream, ReadableStream, toStream, toReadable, transform, consumeStreamWithLimit } from 'vs/base/common/stream';
7+
import { isReadableStream, newWriteableStream, Readable, consumeReadable, peekReadable, consumeStream, ReadableStream, toStream, toReadable, transform, peekStream, isReadableBufferedStream } from 'vs/base/common/stream';
88
import { timeout } from 'vs/base/common/async';
99

1010
suite('Stream', () => {
@@ -14,6 +14,15 @@ suite('Stream', () => {
1414
assert.ok(isReadableStream(newWriteableStream(d => d)));
1515
});
1616

17+
test('isReadableBufferedStream', async () => {
18+
assert.ok(!isReadableBufferedStream(Object.create(null)));
19+
20+
const stream = newWriteableStream(d => d);
21+
stream.end();
22+
const bufferedStream = await peekStream(stream, 1);
23+
assert.ok(isReadableBufferedStream(bufferedStream));
24+
});
25+
1726
test('WriteableStream - basics', () => {
1827
const stream = newWriteableStream<string>(strings => strings.join());
1928

@@ -148,11 +157,11 @@ suite('Stream', () => {
148157
assert.equal(consumed, '1,2,3,4,5');
149158
});
150159

151-
test('consumeReadableWithLimit', () => {
160+
test('peekReadable', () => {
152161
for (let i = 0; i < 5; i++) {
153162
const readable = arrayToReadable(['1', '2', '3', '4', '5']);
154163

155-
const consumedOrReadable = consumeReadableWithLimit(readable, strings => strings.join(), i);
164+
const consumedOrReadable = peekReadable(readable, strings => strings.join(), i);
156165
if (typeof consumedOrReadable === 'string') {
157166
assert.fail('Unexpected result');
158167
} else {
@@ -162,11 +171,11 @@ suite('Stream', () => {
162171
}
163172

164173
let readable = arrayToReadable(['1', '2', '3', '4', '5']);
165-
let consumedOrReadable = consumeReadableWithLimit(readable, strings => strings.join(), 5);
174+
let consumedOrReadable = peekReadable(readable, strings => strings.join(), 5);
166175
assert.equal(consumedOrReadable, '1,2,3,4,5');
167176

168177
readable = arrayToReadable(['1', '2', '3', '4', '5']);
169-
consumedOrReadable = consumeReadableWithLimit(readable, strings => strings.join(), 6);
178+
consumedOrReadable = peekReadable(readable, strings => strings.join(), 6);
170179
assert.equal(consumedOrReadable, '1,2,3,4,5');
171180
});
172181

@@ -198,26 +207,39 @@ suite('Stream', () => {
198207
assert.equal(consumed, '1,2,3,4,5');
199208
});
200209

201-
test('consumeStreamWithLimit', async () => {
210+
test('peekStream', async () => {
202211
for (let i = 0; i < 5; i++) {
203-
const readable = readableToStream(arrayToReadable(['1', '2', '3', '4', '5']));
212+
const stream = readableToStream(arrayToReadable(['1', '2', '3', '4', '5']));
204213

205-
const consumedOrStream = await consumeStreamWithLimit(readable, strings => strings.join(), i);
206-
if (typeof consumedOrStream === 'string') {
207-
assert.fail('Unexpected result');
214+
const result = await peekStream(stream, i);
215+
assert.equal(stream, result.stream);
216+
if (result.ended) {
217+
assert.fail('Unexpected result, stream should not have ended yet');
208218
} else {
209-
const consumed = await consumeStream(consumedOrStream, strings => strings.join());
210-
assert.equal(consumed, '1,2,3,4,5');
219+
assert.equal(result.buffer.length, i + 1, `maxChunks: ${i}`);
220+
221+
const additionalResult: string[] = [];
222+
await consumeStream(stream, strings => {
223+
additionalResult.push(...strings);
224+
225+
return strings.join();
226+
});
227+
228+
assert.equal([...result.buffer, ...additionalResult].join(), '1,2,3,4,5');
211229
}
212230
}
213231

214232
let stream = readableToStream(arrayToReadable(['1', '2', '3', '4', '5']));
215-
let consumedOrStream = await consumeStreamWithLimit(stream, strings => strings.join(), 5);
216-
assert.equal(consumedOrStream, '1,2,3,4,5');
233+
let result = await peekStream(stream, 5);
234+
assert.equal(stream, result.stream);
235+
assert.equal(result.buffer.join(), '1,2,3,4,5');
236+
assert.equal(result.ended, true);
217237

218238
stream = readableToStream(arrayToReadable(['1', '2', '3', '4', '5']));
219-
consumedOrStream = await consumeStreamWithLimit(stream, strings => strings.join(), 6);
220-
assert.equal(consumedOrStream, '1,2,3,4,5');
239+
result = await peekStream(stream, 6);
240+
assert.equal(stream, result.stream);
241+
assert.equal(result.buffer.join(), '1,2,3,4,5');
242+
assert.equal(result.ended, true);
221243
});
222244

223245
test('toStream', async () => {

‎src/vs/base/test/node/buffer.test.ts‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,9 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import * as assert from 'assert';
7-
import { VSBuffer, bufferToReadable, readableToBuffer, bufferToStream, streamToBuffer, newWriteableBufferStream } from 'vs/base/common/buffer';
7+
import { VSBuffer, bufferToReadable, readableToBuffer, bufferToStream, streamToBuffer, newWriteableBufferStream, bufferedStreamToBuffer } from 'vs/base/common/buffer';
88
import { timeout } from 'vs/base/common/async';
9+
import { peekStream } from 'vs/base/common/stream';
910

1011
suite('Buffer', () => {
1112

@@ -29,6 +30,13 @@ suite('Buffer', () => {
2930
assert.equal((await streamToBuffer(stream)).toString(), content);
3031
});
3132

33+
test('bufferedStreamToBuffer', async () => {
34+
const content = 'Hello World';
35+
const stream = await peekStream(bufferToStream(VSBuffer.fromString(content)), 1);
36+
37+
assert.equal((await bufferedStreamToBuffer(stream)).toString(), content);
38+
});
39+
3240
test('bufferWriteableStream - basics (no error)', async () => {
3341
const stream = newWriteableBufferStream();
3442

0 commit comments

Comments
 (0)