Skip to content

Commit 0c73b69

Browse files
author
Benjamin Pasero
committed
streams - add highWaterMark option
1 parent a298366 commit 0c73b69

4 files changed

Lines changed: 156 additions & 9 deletions

File tree

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -214,6 +214,6 @@ export function streamToBufferReadableStream(stream: streams.ReadableStreamEvent
214214
return streams.transform<Uint8Array | string, VSBuffer>(stream, { data: data => typeof data === 'string' ? VSBuffer.fromString(data) : VSBuffer.wrap(data) }, chunks => VSBuffer.concat(chunks));
215215
}
216216

217-
export function newWriteableBufferStream(): streams.WriteableStream<VSBuffer> {
218-
return streams.newWriteableStream<VSBuffer>(chunks => VSBuffer.concat(chunks));
217+
export function newWriteableBufferStream(options?: streams.WriteableStreamOptions): streams.WriteableStream<VSBuffer> {
218+
return streams.newWriteableStream<VSBuffer>(chunks => VSBuffer.concat(chunks), options);
219219
}

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

Lines changed: 69 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,11 @@ export interface ReadableStream<T> extends ReadableStreamEvents<T> {
4949
* Destroys the stream and stops emitting any event.
5050
*/
5151
destroy(): void;
52+
53+
/**
54+
* Allows to remove a listener that was previously added.
55+
*/
56+
removeListener(event: string, callback: Function): void;
5257
}
5358

5459
/**
@@ -74,8 +79,14 @@ export interface WriteableStream<T> extends ReadableStream<T> {
7479
* Writing data to the stream will trigger the on('data')
7580
* event listener if the stream is flowing and buffer the
7681
* data otherwise until the stream is flowing.
82+
*
83+
* If a `highWaterMark` is configured and writing to the
84+
* stream reaches this mark, a promise will be returned
85+
* that should be awaited on before writing more data.
86+
* Otherwise there is a risk of buffering a large number
87+
* of data chunks without consumer.
7788
*/
78-
write(data: T): void;
89+
write(data: T): void | Promise<void>;
7990

8091
/**
8192
* Signals an error to the consumer of the stream via the
@@ -118,8 +129,18 @@ export interface ITransformer<Original, Transformed> {
118129
error?: IErrorTransformer;
119130
}
120131

121-
export function newWriteableStream<T>(reducer: IReducer<T>): WriteableStream<T> {
122-
return new WriteableStreamImpl<T>(reducer);
132+
export function newWriteableStream<T>(reducer: IReducer<T>, options?: WriteableStreamOptions): WriteableStream<T> {
133+
return new WriteableStreamImpl<T>(reducer, options);
134+
}
135+
136+
export interface WriteableStreamOptions {
137+
138+
/**
139+
* The number of objects to buffer before WriteableStream#write()
140+
* signals back that the buffer is full. Can be used to reduce
141+
* the memory pressure when the stream is not flowing.
142+
*/
143+
highWaterMark?: number;
123144
}
124145

125146
class WriteableStreamImpl<T> implements WriteableStream<T> {
@@ -141,7 +162,9 @@ class WriteableStreamImpl<T> implements WriteableStream<T> {
141162
end: [] as { (): void }[]
142163
};
143164

144-
constructor(private reducer: IReducer<T>) { }
165+
private readonly pendingWritePromises: Function[] = [];
166+
167+
constructor(private reducer: IReducer<T>, private options?: WriteableStreamOptions) { }
145168

146169
pause(): void {
147170
if (this.state.destroyed) {
@@ -166,7 +189,7 @@ class WriteableStreamImpl<T> implements WriteableStream<T> {
166189
}
167190
}
168191

169-
write(data: T): void {
192+
write(data: T): void | Promise<void> {
170193
if (this.state.destroyed) {
171194
return;
172195
}
@@ -179,6 +202,11 @@ class WriteableStreamImpl<T> implements WriteableStream<T> {
179202
// not yet flowing: buffer data until flowing
180203
else {
181204
this.buffer.data.push(data);
205+
206+
// highWaterMark: if configured, signal back when buffer reached limits
207+
if (typeof this.options?.highWaterMark === 'number' && this.buffer.data.length > this.options.highWaterMark) {
208+
return new Promise(resolve => this.pendingWritePromises.push(resolve));
209+
}
182210
}
183211
}
184212

@@ -267,13 +295,47 @@ class WriteableStreamImpl<T> implements WriteableStream<T> {
267295
}
268296
}
269297

298+
removeListener(event: string, callback: Function): void {
299+
if (this.state.destroyed) {
300+
return;
301+
}
302+
303+
let listeners: unknown[] | undefined = undefined;
304+
305+
switch (event) {
306+
case 'data':
307+
listeners = this.listeners.data;
308+
break;
309+
310+
case 'end':
311+
listeners = this.listeners.end;
312+
break;
313+
314+
case 'error':
315+
listeners = this.listeners.error;
316+
break;
317+
}
318+
319+
if (listeners) {
320+
const index = listeners.indexOf(callback);
321+
if (index >= 0) {
322+
listeners.splice(index, 1);
323+
}
324+
}
325+
}
326+
270327
private flowData(): void {
271328
if (this.buffer.data.length > 0) {
272329
const fullDataBuffer = this.reducer(this.buffer.data);
273330

274331
this.listeners.data.forEach(listener => listener(fullDataBuffer));
275332

276333
this.buffer.data.length = 0;
334+
335+
// When the buffer is empty, resolve all pending writers
336+
const pendingWritePromises = [...this.pendingWritePromises];
337+
this.pendingWritePromises.length = 0;
338+
pendingWritePromises.forEach(pendingWritePromise => pendingWritePromise());
277339
}
278340
}
279341

@@ -308,6 +370,8 @@ class WriteableStreamImpl<T> implements WriteableStream<T> {
308370
this.listeners.data.length = 0;
309371
this.listeners.error.length = 0;
310372
this.listeners.end.length = 0;
373+
374+
this.pendingWritePromises.length = 0;
311375
}
312376
}
313377
}

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

Lines changed: 83 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55

66
import * as assert from 'assert';
77
import { isReadableStream, newWriteableStream, Readable, consumeReadable, consumeReadableWithLimit, consumeStream, ReadableStream, toStream, toReadable, transform, consumeStreamWithLimit } from 'vs/base/common/stream';
8+
import { timeout } from 'vs/base/common/async';
89

910
suite('Stream', () => {
1011

@@ -13,7 +14,7 @@ suite('Stream', () => {
1314
assert.ok(isReadableStream(newWriteableStream(d => d)));
1415
});
1516

16-
test('WriteableStream', () => {
17+
test('WriteableStream - basics', () => {
1718
const stream = newWriteableStream<string>(strings => strings.join());
1819

1920
let error = false;
@@ -66,6 +67,87 @@ suite('Stream', () => {
6667
assert.equal(chunks.length, 4);
6768
});
6869

70+
test('WriteableStream - removeListener', () => {
71+
const stream = newWriteableStream<string>(strings => strings.join());
72+
73+
let error = false;
74+
const errorListener = (e: Error) => {
75+
error = true;
76+
};
77+
stream.on('error', errorListener);
78+
79+
let end = false;
80+
const endListener = () => {
81+
end = true;
82+
};
83+
stream.on('end', endListener);
84+
85+
let data = false;
86+
const dataListener = () => {
87+
data = true;
88+
};
89+
stream.on('data', dataListener);
90+
91+
stream.write('Hello');
92+
assert.equal(data, true);
93+
94+
data = false;
95+
stream.removeListener('data', dataListener);
96+
97+
stream.write('World');
98+
assert.equal(data, false);
99+
100+
stream.error(new Error());
101+
assert.equal(error, true);
102+
103+
error = false;
104+
stream.removeListener('error', errorListener);
105+
106+
stream.error(new Error());
107+
assert.equal(error, false);
108+
});
109+
110+
test('WriteableStream - highWaterMark', async () => {
111+
const stream = newWriteableStream<string>(strings => strings.join(), { highWaterMark: 3 });
112+
113+
let res = stream.write('1');
114+
assert.ok(!res);
115+
116+
res = stream.write('2');
117+
assert.ok(!res);
118+
119+
res = stream.write('3');
120+
assert.ok(!res);
121+
122+
let promise1 = stream.write('4');
123+
assert.ok(promise1 instanceof Promise);
124+
125+
let promise2 = stream.write('5');
126+
assert.ok(promise2 instanceof Promise);
127+
128+
let drained1 = false;
129+
(async () => {
130+
await promise1;
131+
drained1 = true;
132+
})();
133+
134+
let drained2 = false;
135+
(async () => {
136+
await promise2;
137+
drained2 = true;
138+
})();
139+
140+
let data: string | undefined = undefined;
141+
stream.on('data', chunk => {
142+
data = chunk;
143+
});
144+
assert.ok(data);
145+
146+
await timeout(0);
147+
assert.equal(drained1, true);
148+
assert.equal(drained2, true);
149+
});
150+
69151
test('consumeReadable', () => {
70152
const readable = arrayToReadable(['1', '2', '3', '4', '5']);
71153
const consumed = consumeReadable(readable, strings => strings.join());

‎src/vs/workbench/test/browser/workbenchTestServices.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -752,7 +752,7 @@ export class TestFileService implements IFileService {
752752
this.lastReadFileUri = resource;
753753

754754
return Promise.resolve({
755-
resource: resource,
755+
resource,
756756
value: {
757757
on: (event: string, callback: Function): void => {
758758
if (event === 'data') {
@@ -762,6 +762,7 @@ export class TestFileService implements IFileService {
762762
callback();
763763
}
764764
},
765+
removeListener: () => { },
765766
resume: () => { },
766767
pause: () => { },
767768
destroy: () => { }

0 commit comments

Comments
 (0)