Skip to content

Commit 85c7cc4

Browse files
committed
use promise reject/resolve if udp socket callback has error
1 parent 5fd9b77 commit 85c7cc4

File tree

4 files changed

+104
-58
lines changed

4 files changed

+104
-58
lines changed

aws-distro-opentelemetry-node-autoinstrumentation/src/otlp-udp-exporter.ts

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -24,20 +24,25 @@ export class UdpExporter {
2424
this._socket.unref();
2525
}
2626

27-
sendData(data: Uint8Array, signalFormatPrefix: string): void {
27+
sendData(data: Uint8Array, signalFormatPrefix: string): Promise<void> {
2828
const base64EncodedString = Buffer.from(data).toString('base64');
2929
const message = `${PROTOCOL_HEADER}${signalFormatPrefix}${base64EncodedString}`;
3030

31-
try {
32-
this._socket.send(Buffer.from(message, 'utf-8'), this._port, this._host, err => {
33-
if (err) {
34-
diag.error('Error sending UDP data: %s', err);
35-
}
36-
});
37-
} catch (err) {
38-
diag.error('Error sending UDP data: %s', err);
39-
throw err;
40-
}
31+
return new Promise((resolve, reject) => {
32+
try {
33+
this._socket.send(Buffer.from(message, 'utf-8'), this._port, this._host, err => {
34+
if (err) {
35+
diag.error('Error sending UDP data: %s', err);
36+
reject(err);
37+
} else {
38+
resolve();
39+
}
40+
});
41+
} catch (err) {
42+
diag.error('Error sending UDP data: %s', err);
43+
reject(err);
44+
}
45+
});
4146
}
4247

4348
shutdown(): void {
@@ -70,13 +75,13 @@ export class OTLPUdpSpanExporter implements SpanExporter {
7075
if (serializedData == null) {
7176
return;
7277
}
73-
try {
74-
this._udpExporter.sendData(serializedData, this._signalPrefix);
75-
return resultCallback({ code: ExportResultCode.SUCCESS });
76-
} catch (err) {
77-
diag.error('Error exporting spans: %s', err);
78-
return resultCallback({ code: ExportResultCode.FAILED });
79-
}
78+
this._udpExporter
79+
.sendData(serializedData, this._signalPrefix)
80+
.then(() => resultCallback({ code: ExportResultCode.SUCCESS }))
81+
.catch(err => {
82+
diag.error('Error exporting spans: %s', err);
83+
resultCallback({ code: ExportResultCode.FAILED });
84+
});
8085
}
8186

8287
forceFlush(): Promise<void> {

aws-distro-opentelemetry-node-autoinstrumentation/test/otlp-udp-exporter.test.ts

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -49,18 +49,26 @@ describe('UdpExporterTest', () => {
4949
expect(socketSend.getCall(0).args[0]).toEqual(protbufBinary);
5050
});
5151

52-
it('should handle errors when sending UDP data', () => {
52+
it('should handle errors when sending UDP data', async () => {
5353
const errorMessage = 'UDP send error';
5454
socketSend.yields(new Error(errorMessage)); // Simulate an error
5555

5656
const data = new Uint8Array([1, 2, 3]);
57-
// Expect the sendData method to throw the error
58-
expect(() => udpExporter.sendData(data, 'T1')).toThrow(errorMessage);
57+
// Expect the sendData method to reject with the error
58+
await expect(udpExporter.sendData(data, 'T1')).rejects.toThrow(errorMessage);
5959
// Assert that diag.error was called with the correct error message
6060
expect(diagErrorSpy.calledOnce).toBe(true);
6161
expect(diagErrorSpy.calledWith('Error sending UDP data: %s', sinon.match.instanceOf(Error))).toBe(true);
6262
});
6363

64+
it('should resolve when sending UDP data successfully', async () => {
65+
socketSend.yields(null); // Simulate successful send
66+
67+
const data = new Uint8Array([1, 2, 3]);
68+
await expect(udpExporter.sendData(data, 'T1')).resolves.toBeUndefined();
69+
expect(diagErrorSpy.notCalled).toBe(true);
70+
});
71+
6472
it('should close the socket on shutdown', () => {
6573
udpExporter.shutdown();
6674
expect(socketClose.calledOnce).toBe(true);
@@ -125,16 +133,21 @@ describe('OTLPUdpSpanExporterTest', () => {
125133
sinon.restore();
126134
});
127135

128-
it('should export spans successfully', () => {
136+
it('should export spans successfully', done => {
129137
const callback = sinon.stub();
130138
// Stub ProtobufTraceSerializer.serializeRequest
131139
sinon.stub(ProtobufTraceSerializer, 'serializeRequest').returns(serializedData);
140+
udpExporterMock.sendData.resolves(); // Make sendData resolve successfully
132141

133142
otlpUdpSpanExporter.export(spans, callback);
134143

135-
expect(udpExporterMock.sendData.calledOnceWith(serializedData, 'T1')).toBe(true);
136-
expect(callback.calledOnceWith({ code: ExportResultCode.SUCCESS })).toBe(true);
137-
expect(diagErrorSpy.notCalled).toBe(true); // Ensure no error was logged
144+
// Use setTimeout to allow Promise to resolve
145+
setTimeout(() => {
146+
expect(udpExporterMock.sendData.calledOnceWith(serializedData, 'T1')).toBe(true);
147+
expect(callback.calledOnceWith({ code: ExportResultCode.SUCCESS })).toBe(true);
148+
expect(diagErrorSpy.notCalled).toBe(true); // Ensure no error was logged
149+
done();
150+
}, 0);
138151
});
139152

140153
it('should handle serialization failure', () => {
@@ -149,16 +162,21 @@ describe('OTLPUdpSpanExporterTest', () => {
149162
expect(diagErrorSpy.notCalled).toBe(true);
150163
});
151164

152-
it('should handle errors during export', () => {
165+
it('should handle errors during export', done => {
153166
const error = new Error('Export error');
154-
udpExporterMock.sendData.throws(error);
167+
udpExporterMock.sendData.rejects(error); // Make sendData reject with error
168+
sinon.stub(ProtobufTraceSerializer, 'serializeRequest').returns(serializedData);
155169

156170
const callback = sinon.stub();
157171

158172
otlpUdpSpanExporter.export(spans, callback);
159173

160-
expect(diagErrorSpy.calledOnceWith('Error exporting spans: %s', sinon.match.instanceOf(Error))).toBe(true);
161-
expect(callback.calledOnceWith({ code: ExportResultCode.FAILED })).toBe(true);
174+
// Use setTimeout to allow Promise to reject
175+
setTimeout(() => {
176+
expect(diagErrorSpy.calledOnceWith('Error exporting spans: %s', error)).toBe(true);
177+
expect(callback.calledOnceWith({ code: ExportResultCode.FAILED })).toBe(true);
178+
done();
179+
}, 0);
162180
});
163181

164182
it('should shutdown the UDP exporter successfully', async () => {

exporters/aws-distro-opentelemetry-exporter-xray-udp/src/aws-xray-udp-span-exporter.ts

Lines changed: 23 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -24,20 +24,25 @@ export class UdpExporter {
2424
this._socket.unref();
2525
}
2626

27-
sendData(data: Uint8Array, signalFormatPrefix: string): void {
27+
sendData(data: Uint8Array, signalFormatPrefix: string): Promise<void> {
2828
const base64EncodedString = Buffer.from(data).toString('base64');
2929
const message = `${PROTOCOL_HEADER}${signalFormatPrefix}${base64EncodedString}`;
3030

31-
try {
32-
this._socket.send(Buffer.from(message, 'utf-8'), this._port, this._host, err => {
33-
if (err) {
34-
diag.error('Error sending UDP data: %s', err);
35-
}
36-
});
37-
} catch (err) {
38-
diag.error('Error sending UDP data: %s', err);
39-
throw err;
40-
}
31+
return new Promise((resolve, reject) => {
32+
try {
33+
this._socket.send(Buffer.from(message, 'utf-8'), this._port, this._host, err => {
34+
if (err) {
35+
diag.error('Error sending UDP data: %s', err);
36+
reject(err);
37+
} else {
38+
resolve();
39+
}
40+
});
41+
} catch (err) {
42+
diag.error('Error sending UDP data: %s', err);
43+
reject(err);
44+
}
45+
});
4146
}
4247

4348
shutdown(): void {
@@ -79,13 +84,13 @@ export class AwsXrayUdpSpanExporter implements SpanExporter {
7984
if (serializedData == null) {
8085
return;
8186
}
82-
try {
83-
this._udpExporter.sendData(serializedData, this._signalPrefix);
84-
return resultCallback({ code: ExportResultCode.SUCCESS });
85-
} catch (err) {
86-
diag.error('Error exporting spans: %s', err);
87-
return resultCallback({ code: ExportResultCode.FAILED });
88-
}
87+
this._udpExporter
88+
.sendData(serializedData, this._signalPrefix)
89+
.then(() => resultCallback({ code: ExportResultCode.SUCCESS }))
90+
.catch(err => {
91+
diag.error('Error exporting spans: %s', err);
92+
resultCallback({ code: ExportResultCode.FAILED });
93+
});
8994
}
9095

9196
forceFlush(): Promise<void> {

exporters/aws-distro-opentelemetry-exporter-xray-udp/test/aws-xray-udp-span-exporter.test.ts

Lines changed: 29 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -50,18 +50,26 @@ describe('UdpExporterTest', () => {
5050
expect(socketSend.getCall(0).args[0]).toEqual(protbufBinary);
5151
});
5252

53-
it('should handle errors when sending UDP data', () => {
53+
it('should handle errors when sending UDP data', async () => {
5454
const errorMessage = 'UDP send error';
5555
socketSend.yields(new Error(errorMessage)); // Simulate an error
5656

5757
const data = new Uint8Array([1, 2, 3]);
58-
// Expect the sendData method to throw the error
59-
expect(() => udpExporter.sendData(data, 'T1')).toThrow(errorMessage);
58+
// Expect the sendData method to reject with the error
59+
await expect(udpExporter.sendData(data, 'T1')).rejects.toThrow(errorMessage);
6060
// Assert that diag.error was called with the correct error message
6161
expect(diagErrorSpy.calledOnce).toBe(true);
6262
expect(diagErrorSpy.calledWith('Error sending UDP data: %s', sinon.match.instanceOf(Error))).toBe(true);
6363
});
6464

65+
it('should resolve when sending UDP data successfully', async () => {
66+
socketSend.yields(null); // Simulate successful send
67+
68+
const data = new Uint8Array([1, 2, 3]);
69+
await expect(udpExporter.sendData(data, 'T1')).resolves.toBeUndefined();
70+
expect(diagErrorSpy.notCalled).toBe(true);
71+
});
72+
6573
it('should close the socket on shutdown', () => {
6674
udpExporter.shutdown();
6775
expect(socketClose.calledOnce).toBe(true);
@@ -132,16 +140,21 @@ describe('AwsXrayUdpSpanExporterTest', () => {
132140
sinon.restore();
133141
});
134142

135-
it('should export spans successfully', () => {
143+
it('should export spans successfully', done => {
136144
const callback = sinon.stub();
137145
// Stub ProtobufTraceSerializer.serializeRequest
138146
sinon.stub(ProtobufTraceSerializer, 'serializeRequest').returns(serializedData);
147+
udpExporterMock.sendData.resolves(); // Make sendData resolve successfully
139148

140149
awsXrayUdpSpanExporter.export(spans, callback);
141150

142-
expect(udpExporterMock.sendData.calledOnceWith(serializedData, 'T1')).toBe(true);
143-
expect(callback.calledOnceWith({ code: ExportResultCode.SUCCESS })).toBe(true);
144-
expect(diagErrorSpy.notCalled).toBe(true); // Ensure no error was logged
151+
// Use setTimeout to allow Promise to resolve
152+
setTimeout(() => {
153+
expect(udpExporterMock.sendData.calledOnceWith(serializedData, 'T1')).toBe(true);
154+
expect(callback.calledOnceWith({ code: ExportResultCode.SUCCESS })).toBe(true);
155+
expect(diagErrorSpy.notCalled).toBe(true); // Ensure no error was logged
156+
done();
157+
}, 0);
145158
});
146159

147160
it('should handle serialization failure', () => {
@@ -156,16 +169,21 @@ describe('AwsXrayUdpSpanExporterTest', () => {
156169
expect(diagErrorSpy.notCalled).toBe(true);
157170
});
158171

159-
it('should handle errors during export', () => {
172+
it('should handle errors during export', done => {
160173
const error = new Error('Export error');
161-
udpExporterMock.sendData.throws(error);
174+
udpExporterMock.sendData.rejects(error);
175+
sinon.stub(ProtobufTraceSerializer, 'serializeRequest').returns(serializedData);
162176

163177
const callback = sinon.stub();
164178

165179
awsXrayUdpSpanExporter.export(spans, callback);
166180

167-
expect(diagErrorSpy.calledOnceWith('Error exporting spans: %s', sinon.match.instanceOf(Error))).toBe(true);
168-
expect(callback.calledOnceWith({ code: ExportResultCode.FAILED })).toBe(true);
181+
// Use setTimeout to allow Promise to reject
182+
setTimeout(() => {
183+
expect(diagErrorSpy.calledOnceWith('Error exporting spans: %s', error)).toBe(true);
184+
expect(callback.calledOnceWith({ code: ExportResultCode.FAILED })).toBe(true);
185+
done();
186+
}, 0);
169187
});
170188

171189
it('should forceFlush without throwing', async () => {

0 commit comments

Comments
 (0)