update to xray segments

leighton-digital · Aug 29, 2024 · 0bcf6ba · 0bcf6ba
1 parent 3e03c3c
commit 0bcf6ba
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/lib/cache/redisClient.js b/lib/cache/redisClient.js
@@ -18,6 +18,7 @@ const sts = new AWS.STS({
 
 async function assumeRole() {
   const segment = AWSXRay.getSegment() || new AWSXRay.Segment('assumeRole');
+  AWSXRay.setSegment(segment)
   try{
     const token = fs.readFileSync(tokenFile, 'utf8');
 

diff --git a/lib/mq/mqUtils.js b/lib/mq/mqUtils.js
@@ -23,7 +23,7 @@ async function connect() {
   const segment = process.env.NODE_ENV === 'development'
     ? new MockSegment('RabbitMQ-Connect')
     : AWSXRay.getSegment() || new AWSXRay.Segment('RabbitMQ-Connect');
-
+    AWSXRay.setSegment(segment)
   try {
     const connection = await amqp.connect(RABBITMQ_URL);
     connection.on('error', (error) => {
@@ -52,7 +52,7 @@ const retryConnect = async (maxRetries = 5, delay = 1000) => {
   const segment = process.env.NODE_ENV === 'development'
     ? new MockSegment('RabbitMQ-RetryConnect')
     : AWSXRay.getSegment() || new AWSXRay.Segment('RabbitMQ-RetryConnect');
-
+    AWSXRay.setSegment(segment)
   for (let i = 0; i < maxRetries; i++) {
     try {
       const connectionData = await connect();
@@ -75,7 +75,7 @@ async function createAndBindQueue(channel) {
   const segment = process.env.NODE_ENV === 'development'
     ? new MockSegment('RabbitMQ-CreateAndBindQueue')
     : AWSXRay.getSegment() || new AWSXRay.Segment('RabbitMQ-CreateAndBindQueue');
-
+  AWSXRay.setSegment(segment)
   try {
     logger.info(`Using queue name: ${QUEUE_NAME}`);
     await channel.assertExchange('fanoutExchange', 'fanout'); 
@@ -96,7 +96,7 @@ async function listenToQueue(channel, callback) {
   const segment = process.env.NODE_ENV === 'development'
   ? new MockSegment('RabbitMQ-ListenToQueue')
   : AWSXRay.getSegment() || new AWSXRay.Segment('RabbitMQ-ListenToQueue');
-
+  AWSXRay.setSegment(segment)
   channel.consume(QUEUE_NAME, (message) => {
     const subsegment = segment.addNewSubsegment('Process-Message');
     try {
@@ -119,7 +119,7 @@ async function publishMessage(channel, message) {
   const segment = process.env.NODE_ENV === 'development'
   ? new MockSegment('RabbitMQ-PublishMessage')
   : AWSXRay.getSegment() || new AWSXRay.Segment('RabbitMQ-PublishMessage');
-
+  AWSXRay.setSegment(segment)
   try {
     channel.publish('fanoutExchange', '', Buffer.from(JSON.stringify(message)));
   } catch (publishError) {