Skip to content

Commit a8b1b81

Browse files
petermyerspetermyers
authored
Merge pull request #410 from jgilbert01/issue-sqs-batching-by-size
Batch by total request message size.
2 parents c7403f6 + 3044336 commit a8b1b81

6 files changed

Lines changed: 419 additions & 6 deletions

File tree

‎package-lock.json‎

Lines changed: 2 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎package.json‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name": "aws-lambda-stream",
3-
"version": "1.1.8",
3+
"version": "1.1.9",
44
"description": "Create stream processors with AWS Lambda functions.",
55
"keywords": [
66
"aws",

‎src/sinks/sqs.js‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@ import _ from 'highland';
22

33
import Connector from '../connectors/sqs';
44

5-
import { toBatchUow, unBatchUow } from '../utils/batch';
5+
import { batchWithPayloadSizeOrCount, toBatchUow, unBatchUow } from '../utils/batch';
66
import { ratelimit } from '../utils/ratelimit';
77
import { rejectWithFault } from '../utils/faults';
88
import { debug as d } from '../utils/print';
@@ -13,6 +13,7 @@ export const sendToSqs = ({ // eslint-disable-line import/prefer-default-export
1313
queueUrl = process.env.QUEUE_URL,
1414
messageField = 'message',
1515
batchSize = Number(process.env.SQS_BATCH_SIZE) || Number(process.env.BATCH_SIZE) || 10,
16+
maxPayloadSize = Number(process.env.SQS_MAX_PAYLOAD_SIZE) || Number(process.env.MAX_PAYLOAD_SIZE) || 1024 * 1024,
1617
parallel = Number(process.env.SQS_PARALLEL) || Number(process.env.PARALLEL) || 8,
1718
step = 'send',
1819
...opt
@@ -46,7 +47,12 @@ export const sendToSqs = ({ // eslint-disable-line import/prefer-default-export
4647
return (s) => s
4748
.through(ratelimit(opt))
4849

49-
.batch(batchSize)
50+
.consume(batchWithPayloadSizeOrCount({
51+
batchSize,
52+
maxPayloadSize,
53+
payloadField: messageField,
54+
...opt,
55+
}))
5056
.map(toBatchUow)
5157

5258
.map(toInputParams)

‎src/utils/batch.js‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import _ from 'highland';
22
import isFunction from 'lodash/isFunction';
3+
import { get } from 'lodash';
34
import { toClaimcheckEvent, toPutClaimcheckRequest } from '../sinks/claimcheck';
45

56
// used after highland batch step
@@ -48,6 +49,9 @@ export const compact = (rule) => {
4849
})));
4950
};
5051

52+
/**
53+
* Batch EB request entries by size to avoid writing a batch that's too large to EB.
54+
*/
5155
export const batchWithSize = ({
5256
claimCheckBucketName = process.env.CLAIMCHECK_BUCKET_NAME,
5357
putClaimcheckRequest = 'putClaimcheckRequest',
@@ -116,3 +120,58 @@ const logMetrics = (batch, sizes, opt) => {
116120
batch[0].metrics?.gauge('publish|stream.pipeline.eventSize.bytes', sizes);
117121
}
118122
};
123+
124+
/**
125+
* Batches by aggregate payload size with a cap on payload count.
126+
*/
127+
export const batchWithPayloadSizeOrCount = ({
128+
batchSize,
129+
maxPayloadSize,
130+
payloadField,
131+
...opt
132+
}) => {
133+
let batched = [];
134+
let sizes = [];
135+
136+
return (err, x, push, next) => {
137+
/* istanbul ignore if */
138+
if (err) {
139+
push(err);
140+
next();
141+
} else if (x === _.nil) {
142+
if (batched.length > 0) {
143+
logMetrics(batched, sizes, opt);
144+
push(null, batched);
145+
}
146+
147+
push(null, _.nil);
148+
} else {
149+
if (!get(x, payloadField)) {
150+
push(null, [x]);
151+
} else {
152+
const size = Buffer.byteLength(JSON.stringify(get(x, payloadField)));
153+
if (size > maxPayloadSize) {
154+
logMetrics([x], [size], opt);
155+
const error = new Error(`Payload size: ${size}, exceeded max: ${maxPayloadSize}`);
156+
error.uow = x;
157+
push(error);
158+
next();
159+
return;
160+
}
161+
162+
const totalSize = sizes.reduce((a, c) => a + c, size);
163+
if (totalSize <= maxPayloadSize && batched.length + 1 <= batchSize) {
164+
batched.push(x);
165+
sizes.push(size);
166+
} else {
167+
logMetrics(batched, sizes, opt);
168+
push(null, batched);
169+
batched = [x];
170+
sizes = [size];
171+
}
172+
}
173+
174+
next();
175+
}
176+
};
177+
};

‎test/unit/sinks/sqs.test.js‎

Lines changed: 134 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,11 @@ describe('sinks/sqs.js', () => {
1818
Id: '1',
1919
MessageBody: JSON.stringify({ f1: 'v1' }),
2020
},
21+
}, {
22+
message: {
23+
Id: '2',
24+
MessageBody: JSON.stringify({ f1: 'v2' }),
25+
},
2126
}];
2227

2328
_(uows)
@@ -26,7 +31,68 @@ describe('sinks/sqs.js', () => {
2631
.tap((collected) => {
2732
// console.log(JSON.stringify(collected, null, 2));
2833

29-
expect(collected.length).to.equal(1);
34+
expect(collected.length).to.equal(2);
35+
expect(collected[0]).to.deep.equal({
36+
message: {
37+
Id: '1',
38+
MessageBody: JSON.stringify({ f1: 'v1' }),
39+
},
40+
inputParams: {
41+
Entries: [{
42+
Id: '1',
43+
MessageBody: JSON.stringify({ f1: 'v1' }),
44+
}, {
45+
Id: '2',
46+
MessageBody: JSON.stringify({ f1: 'v2' }),
47+
}],
48+
},
49+
sendMessageBatchResponse: {},
50+
});
51+
52+
expect(collected[1]).to.deep.equal({
53+
message: {
54+
Id: '2',
55+
MessageBody: JSON.stringify({ f1: 'v2' }),
56+
},
57+
inputParams: {
58+
Entries: [{
59+
Id: '1',
60+
MessageBody: JSON.stringify({ f1: 'v1' }),
61+
}, {
62+
Id: '2',
63+
MessageBody: JSON.stringify({ f1: 'v2' }),
64+
}],
65+
},
66+
sendMessageBatchResponse: {},
67+
});
68+
})
69+
.done(done);
70+
});
71+
72+
it('should split a batch due to batch size', (done) => {
73+
sinon.stub(Connector.prototype, 'sendMessageBatch').resolves({});
74+
75+
const uows = [{
76+
message: {
77+
Id: '1',
78+
MessageBody: JSON.stringify({ f1: 'v1' }),
79+
},
80+
}, {
81+
message: {
82+
Id: '2',
83+
MessageBody: JSON.stringify({ f1: 'v2' }),
84+
},
85+
}];
86+
87+
_(uows)
88+
.through(sendToSqs({
89+
batchSize: 1,
90+
}))
91+
.collect()
92+
.tap((collected) => {
93+
// console.log(JSON.stringify(collected, null, 2));
94+
95+
expect(collected.length).to.equal(2);
3096
expect(collected[0]).to.deep.equal({
3197
message: {
3298
Id: '1',
@@ -40,6 +106,73 @@ describe('sinks/sqs.js', () => {
40106
},
41107
sendMessageBatchResponse: {},
42108
});
109+
expect(collected[1]).to.deep.equal({
110+
message: {
111+
Id: '2',
112+
MessageBody: JSON.stringify({ f1: 'v2' }),
113+
},
114+
inputParams: {
115+
Entries: [{
116+
Id: '2',
117+
MessageBody: JSON.stringify({ f1: 'v2' }),
118+
}],
119+
},
120+
sendMessageBatchResponse: {},
121+
});
122+
})
123+
.done(done);
124+
});
125+
126+
it('should split a batch due to payload size', (done) => {
127+
sinon.stub(Connector.prototype, 'sendMessageBatch').resolves({});
128+
129+
const uows = [{
130+
message: {
131+
Id: '1',
132+
MessageBody: JSON.stringify({ f1: 'v1' }),
133+
},
134+
}, {
135+
message: {
136+
Id: '2',
137+
MessageBody: JSON.stringify({ f1: 'v2' }),
138+
},
139+
}];
140+
141+
_(uows)
142+
.through(sendToSqs({
143+
maxPayloadSize: 50,
144+
}))
145+
.collect()
146+
.tap((collected) => {
147+
// console.log(JSON.stringify(collected, null, 2));
148+
149+
expect(collected.length).to.equal(2);
150+
expect(collected[0]).to.deep.equal({
151+
message: {
152+
Id: '1',
153+
MessageBody: JSON.stringify({ f1: 'v1' }),
154+
},
155+
inputParams: {
156+
Entries: [{
157+
Id: '1',
158+
MessageBody: JSON.stringify({ f1: 'v1' }),
159+
}],
160+
},
161+
sendMessageBatchResponse: {},
162+
});
163+
expect(collected[1]).to.deep.equal({
164+
message: {
165+
Id: '2',
166+
MessageBody: JSON.stringify({ f1: 'v2' }),
167+
},
168+
inputParams: {
169+
Entries: [{
170+
Id: '2',
171+
MessageBody: JSON.stringify({ f1: 'v2' }),
172+
}],
173+
},
174+
sendMessageBatchResponse: {},
175+
});
43176
})
44177
.done(done);
45178
});

0 commit comments

Comments
 (0)