Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 508eaef

Browse files
pinzonpinzon
authored
publish SNS messages asynchronously to fix performance issues (#3105)
1 parent 8cc8a3a commit 508eaef

File tree

4 files changed

+118
-70
lines changed

4 files changed

+118
-70
lines changed

‎localstack/services/kinesis/kinesis_listener.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,6 @@ def return_response(self, method, path, data, headers, response):
136136
return response
137137

138138
def decode_content(self, data):
139-
# return json.loads(to_str(data))
140139
try:
141140
return json.loads(to_str(data))
142141
except UnicodeDecodeError:

‎localstack/services/sns/sns_listener.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717
from localstack.utils.aws import aws_stack
1818
from localstack.utils.aws.aws_responses import response_regex_replace
1919
from localstack.utils.aws.dead_letter_queue import sns_error_to_dead_letter_queue
20-
from localstack.utils.common import timestamp_millis, short_uid, to_str, to_bytes
20+
from localstack.utils.common import timestamp_millis, short_uid, to_str, to_bytes, start_thread
2121
from localstack.utils.persistence import PersistingProxyListener
2222

2323
# set up logger
@@ -257,17 +257,7 @@ def unsubscribe_sqs_queue(queue_url):
257257
subscriptions.remove(subscriber)
258258

259259

260-
def publish_message(topic_arn, req_data, subscription_arn=None, skip_checks=False):
261-
message = req_data['Message'][0]
262-
message_id = str(uuid.uuid4())
263-
264-
LOG.debug('Publishing message to TopicArn: %s | Message: %s' % (topic_arn, message))
265-
266-
if topic_arn and ':endpoint/' in topic_arn:
267-
# cache messages published to platform endpoints
268-
cache = PLATFORM_ENDPOINT_MESSAGES[topic_arn] = PLATFORM_ENDPOINT_MESSAGES.get(topic_arn) or []
269-
cache.append(req_data)
270-
260+
def message_to_subscribers(message_id, message, topic_arn, req_data, subscription_arn=None, skip_checks=False):
271261
subscriptions = SNS_SUBSCRIPTIONS.get(topic_arn, [])
272262
for subscriber in list(subscriptions):
273263
if subscription_arn not in [None, subscriber['SubscriptionArn']]:
@@ -375,6 +365,19 @@ def publish_message(topic_arn, req_data, subscription_arn=None, skip_checks=Fals
375365
else:
376366
LOG.warning('Unexpected protocol "%s" for SNS subscription' % subscriber['Protocol'])
377367

368+
369+
def publish_message(topic_arn, req_data, subscription_arn=None, skip_checks=False):
370+
message = req_data['Message'][0]
371+
message_id = str(uuid.uuid4())
372+
373+
if topic_arn and ':endpoint/' in topic_arn:
374+
# cache messages published to platform endpoints
375+
cache = PLATFORM_ENDPOINT_MESSAGES[topic_arn] = PLATFORM_ENDPOINT_MESSAGES.get(topic_arn) or []
376+
cache.append(req_data)
377+
378+
LOG.debug('Publishing message to TopicArn: %s | Message: %s' % (topic_arn, message))
379+
start_thread(
380+
lambda _: message_to_subscribers(message_id, message, topic_arn, req_data, subscription_arn, skip_checks))
378381
return message_id
379382

380383

‎tests/integration/test_notifications.py

Lines changed: 28 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -3,13 +3,15 @@
33
from io import BytesIO
44
from localstack.utils import testutil
55
from localstack.utils.aws import aws_stack
6-
from localstack.utils.common import to_str, short_uid
6+
from localstack.utils.common import to_str, short_uid, retry
77

88
TEST_BUCKET_NAME_WITH_NOTIFICATIONS = 'test-bucket-notif-1'
99
TEST_QUEUE_NAME_FOR_S3 = 'test_queue'
1010
TEST_TOPIC_NAME = 'test_topic_name_for_sqs'
1111
TEST_S3_TOPIC_NAME = 'test_topic_name_for_s3_to_sns_to_sqs'
1212
TEST_QUEUE_NAME_FOR_SNS = 'test_queue_for_sns'
13+
PUBLICATION_TIMEOUT = 0.500
14+
PUBLICATION_RETRIES = 4
1315

1416

1517
class TestNotifications(unittest.TestCase):
@@ -36,14 +38,16 @@ def test_sns_to_sqs(self):
3638
sns_client.publish(TopicArn=topic_info['TopicArn'], Message='test message for SQS',
3739
MessageAttributes={'attr1': {'DataType': 'String', 'StringValue': test_value}})
3840

39-
# receive, assert, and delete message from SQS
40-
queue_url = queue_info['QueueUrl']
41-
assertions = []
42-
# make sure we receive the correct topic ARN in notifications
43-
assertions.append({'TopicArn': topic_info['TopicArn']})
44-
# make sure the notification contains message attributes
45-
assertions.append({'Value': test_value})
46-
self._receive_assert_delete(queue_url, assertions, sqs_client)
41+
def assert_message():
42+
# receive, assert, and delete message from SQS
43+
queue_url = queue_info['QueueUrl']
44+
assertions = []
45+
# make sure we receive the correct topic ARN in notifications
46+
assertions.append({'TopicArn': topic_info['TopicArn']})
47+
# make sure the notification contains message attributes
48+
assertions.append({'Value': test_value})
49+
self._receive_assert_delete(queue_url, assertions, sqs_client)
50+
retry(assert_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
4751

4852
def test_bucket_notifications(self):
4953

@@ -191,20 +195,21 @@ def test_bucket_notifications(self):
191195
s3_client.upload_fileobj(BytesIO(test_data2), TEST_BUCKET_NAME_WITH_NOTIFICATIONS, test_key2)
192196

193197
# verify subject and records
194-
195-
response = sqs_client.receive_message(QueueUrl=queue_url)
196-
for message in response['Messages']:
197-
snsObj = json.loads(message['Body'])
198-
testutil.assert_object({'Subject': 'Amazon S3 Notification'}, snsObj)
199-
notificationObj = json.loads(snsObj['Message'])
200-
testutil.assert_objects(
201-
[
202-
{'key': test_key2},
203-
{'name': TEST_BUCKET_NAME_WITH_NOTIFICATIONS}
204-
], notificationObj['Records'])
205-
206-
sqs_client.delete_message(QueueUrl=queue_url, ReceiptHandle=message['ReceiptHandle'])
207-
198+
def verify():
199+
response = sqs_client.receive_message(QueueUrl=queue_url)
200+
for message in response['Messages']:
201+
snsObj = json.loads(message['Body'])
202+
testutil.assert_object({'Subject': 'Amazon S3 Notification'}, snsObj)
203+
notificationObj = json.loads(snsObj['Message'])
204+
testutil.assert_objects(
205+
[
206+
{'key': test_key2},
207+
{'name': TEST_BUCKET_NAME_WITH_NOTIFICATIONS}
208+
], notificationObj['Records'])
209+
210+
sqs_client.delete_message(QueueUrl=queue_url, ReceiptHandle=message['ReceiptHandle'])
211+
212+
retry(verify, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
208213
self._delete_notification_config()
209214

210215
def _delete_notification_config(self):

‎tests/integration/test_sns.py

Lines changed: 75 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,9 @@
2626
TEST_QUEUE_DLQ_NAME = 'TestQueue_DLQ_snsTest'
2727
TEST_TOPIC_NAME_2 = 'topic-test-2'
2828

29+
PUBLICATION_TIMEOUT = .500
30+
PUBLICATION_RETRIES = 4
31+
2932
THIS_FOLDER = os.path.dirname(os.path.realpath(__file__))
3033
TEST_LAMBDA_ECHO_FILE = os.path.join(THIS_FOLDER, 'lambdas', 'lambda_echo.py')
3134

@@ -54,11 +57,14 @@ def test_publish_unicode_chars(self):
5457
# publish message to SNS, receive it from SQS, assert that messages are equal
5558
message = u'ö§a1"_!?,. £$-'
5659
self.sns_client.publish(TopicArn=self.topic_arn, Message=message)
57-
msgs = self.sqs_client.receive_message(QueueUrl=queue_url)
58-
msg_received = msgs['Messages'][0]
59-
msg_received = json.loads(to_str(msg_received['Body']))
60-
msg_received = msg_received['Message']
61-
self.assertEqual(message, msg_received)
60+
61+
def check_message():
62+
msgs = self.sqs_client.receive_message(QueueUrl=queue_url)
63+
msg_received = msgs['Messages'][0]
64+
msg_received = json.loads(to_str(msg_received['Body']))
65+
msg_received = msg_received['Message']
66+
self.assertEqual(message, msg_received)
67+
retry(check_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
6268

6369
# clean up
6470
self.sqs_client.delete_queue(QueueUrl=queue_url)
@@ -124,11 +130,14 @@ def test_attribute_raw_subscribe(self):
124130
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
125131
MessageAttributes={'store': {'DataType': 'Binary', 'BinaryValue': binary_attribute}})
126132

127-
msgs = self.sqs_client.receive_message(QueueUrl=self.queue_url, MessageAttributeNames=['All'])
128-
msg_received = msgs['Messages'][0]
133+
def check_message():
134+
msgs = self.sqs_client.receive_message(QueueUrl=self.queue_url, MessageAttributeNames=['All'])
135+
msg_received = msgs['Messages'][0]
136+
137+
self.assertEqual(message, msg_received['Body'])
138+
self.assertEqual(binary_attribute, msg_received['MessageAttributes']['store']['BinaryValue'])
129139

130-
self.assertEqual(message, msg_received['Body'])
131-
self.assertEqual(binary_attribute, msg_received['MessageAttributes']['store']['BinaryValue'])
140+
retry(check_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
132141

133142
def test_filter_policy(self):
134143
# connect SNS topic to an SQS queue
@@ -151,15 +160,23 @@ def test_filter_policy(self):
151160
message = u'This is a test message'
152161
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
153162
MessageAttributes={'attr1': {'DataType': 'Number', 'StringValue': '99'}})
154-
num_msgs_1 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
155-
self.assertEqual(num_msgs_1, num_msgs_0 + 1)
163+
164+
def check_message():
165+
num_msgs_1 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
166+
self.assertEqual(num_msgs_1, num_msgs_0 + 1)
167+
return num_msgs_1
168+
num_msgs_1 = retry(check_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
156169

157170
# publish message that does not satisfy the filter policy, assert that message is not received
158171
message = u'This is a test message'
159172
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
160173
MessageAttributes={'attr1': {'DataType': 'Number', 'StringValue': '111'}})
161-
num_msgs_2 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
162-
self.assertEqual(num_msgs_2, num_msgs_1)
174+
175+
def check_message2():
176+
num_msgs_2 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
177+
self.assertEqual(num_msgs_2, num_msgs_1)
178+
return num_msgs_2
179+
retry(check_message2, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
163180

164181
# clean up
165182
self.sqs_client.delete_queue(QueueUrl=queue_url)
@@ -188,15 +205,23 @@ def do_subscribe(self, filter_policy, queue_arn):
188205
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
189206
MessageAttributes={'store': {'DataType': 'Number', 'StringValue': '99'},
190207
'def': {'DataType': 'Number', 'StringValue': '99'}})
191-
num_msgs_1 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
192-
self.assertEqual(num_msgs_1, num_msgs_0 + 1)
208+
209+
def check_message1():
210+
num_msgs_1 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
211+
self.assertEqual(num_msgs_1, num_msgs_0 + 1)
212+
return num_msgs_1
213+
num_msgs_1 = retry(check_message1, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
193214

194215
# publish message that does not satisfy the filter policy, assert that message is not received
195216
message = u'This is a test message'
196217
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
197218
MessageAttributes={'attr1': {'DataType': 'Number', 'StringValue': '111'}})
198-
num_msgs_2 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
199-
self.assertEqual(num_msgs_2, num_msgs_1)
219+
220+
def check_message2():
221+
num_msgs_2 = len(self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages'])
222+
self.assertEqual(num_msgs_2, num_msgs_1)
223+
return num_msgs_2
224+
retry(check_message2, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
200225

201226
# test with exist operator set to false.
202227
queue_arn = aws_stack.sqs_queue_arn(TEST_QUEUE_NAME)
@@ -210,17 +235,27 @@ def do_subscribe(self, filter_policy, queue_arn):
210235
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
211236
MessageAttributes={'store': {'DataType': 'Number', 'StringValue': '99'},
212237
'def': {'DataType': 'Number', 'StringValue': '99'}})
213-
num_msgs_1 = len(self.sqs_client.receive_message(QueueUrl=self.queue_url,
238+
239+
def check_message():
240+
num_msgs_1 = len(self.sqs_client.receive_message(QueueUrl=self.queue_url,
214241
VisibilityTimeout=0).get('Messages', []))
215-
self.assertEqual(num_msgs_1, num_msgs_0)
242+
self.assertEqual(num_msgs_1, num_msgs_0)
243+
return num_msgs_1
244+
245+
num_msgs_1 = retry(check_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
216246

217247
# publish message that without the attribute and see if its getting filtered.
218248
message = u'This is a test message'
219249
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
220250
MessageAttributes={'attr1': {'DataType': 'Number', 'StringValue': '111'}})
221-
num_msgs_2 = len(self.sqs_client.receive_message(QueueUrl=self.queue_url,
222-
VisibilityTimeout=0).get('Messages', []))
223-
self.assertEqual(num_msgs_2, num_msgs_1)
251+
252+
def check_message3():
253+
num_msgs_2 = len(self.sqs_client.receive_message(QueueUrl=self.queue_url,
254+
VisibilityTimeout=0).get('Messages', []))
255+
self.assertEqual(num_msgs_2, num_msgs_1)
256+
return num_msgs_2
257+
258+
retry(check_message3, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
224259

225260
# clean up
226261
self.sqs_client.delete_queue(QueueUrl=queue_url)
@@ -232,8 +267,10 @@ def test_subscribe_sqs_queue(self):
232267
subscription = self._publish_sns_message_with_attrs(queue_arn, 'sqs')
233268

234269
# assert that message is received
235-
messages = self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages']
236-
self.assertEqual(json.loads(messages[0]['Body'])['MessageAttributes']['attr1']['Value'], '99.12')
270+
def check_message():
271+
messages = self.sqs_client.receive_message(QueueUrl=queue_url, VisibilityTimeout=0)['Messages']
272+
self.assertEqual(json.loads(messages[0]['Body'])['MessageAttributes']['attr1']['Value'], '99.12')
273+
retry(check_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
237274

238275
# clean up
239276
self.sqs_client.delete_queue(QueueUrl=queue_url)
@@ -248,7 +285,7 @@ def test_subscribe_platform_endpoint(self):
248285
# assert that message has been received
249286
def check_message():
250287
self.assertGreater(len(sns_listener.PLATFORM_ENDPOINT_MESSAGES[platform_arn]), 0)
251-
retry(check_message)
288+
retry(check_message, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
252289

253290
# clean up
254291
sns.unsubscribe(SubscriptionArn=subscription['SubscriptionArn'])
@@ -270,6 +307,7 @@ def _publish_sns_message_with_attrs(self, endpoint_arn, protocol):
270307
message = u'This is a test message'
271308
self.sns_client.publish(TopicArn=self.topic_arn, Message=message,
272309
MessageAttributes={'attr1': {'DataType': 'Number', 'StringValue': '99.12'}})
310+
time.sleep(PUBLICATION_TIMEOUT)
273311
return subscription
274312

275313
def test_unknown_topic_publish(self):
@@ -355,16 +393,19 @@ def test_topic_subscription(self):
355393
Protocol='email',
356394
Endpoint='[email protected]'
357395
)
358-
subscription_arn = subscription['SubscriptionArn']
359-
subscription_obj = sns_listener.SUBSCRIPTION_STATUS[subscription_arn]
360-
self.assertEqual(subscription_obj['Status'], 'Not Subscribed')
361396

362-
_token = subscription_obj['Token']
363-
self.sns_client.confirm_subscription(
364-
TopicArn=self.topic_arn,
365-
Token=_token
366-
)
367-
self.assertEqual(subscription_obj['Status'], 'Subscribed')
397+
def check_subscription():
398+
subscription_arn = subscription['SubscriptionArn']
399+
subscription_obj = sns_listener.SUBSCRIPTION_STATUS[subscription_arn]
400+
self.assertEqual(subscription_obj['Status'], 'Not Subscribed')
401+
402+
_token = subscription_obj['Token']
403+
self.sns_client.confirm_subscription(
404+
TopicArn=self.topic_arn,
405+
Token=_token
406+
)
407+
self.assertEqual(subscription_obj['Status'], 'Subscribed')
408+
retry(check_subscription, retries=PUBLICATION_RETRIES, sleep=PUBLICATION_TIMEOUT)
368409

369410
def test_dead_letter_queue(self):
370411
lambda_name = 'test-%s' % short_uid()

0 commit comments

Comments
 (0)