Support retry on failure

Author: PetrHeinz

README.md

@@ -62,6 +62,8 @@ The template supports the following optional parameters:
 
 - `batch_size`: Number of messages to batch before sending to Better Stack. Default: 100
 - `window_size`: Window size in seconds for batching messages. Default: 10
+- `max_retries`: Maximum number of retry attempts for failed requests. Default: 3
+- `initial_retry_delay`: Initial delay between retries in seconds. Default: 1
 
 You can include these parameters in your Dataflow job by adding them to the run command, e.g. `gcloud dataflow flex-template run ... --parameters window_size=30`.
 

pipeline.py

@@ -6,6 +6,7 @@
 from apache_beam.transforms.window import FixedWindows
 import requests
 import gzip
+import time
 from typing import Dict, Any, List
 from datetime import timedelta
 
@@ -20,6 +21,8 @@ def __init__(self, source_token: str, ingesting_host: str, batch_size: int):
             'Content-Encoding': 'gzip'
         }
         self.batch = []
+        self.max_retries = 3
+        self.initial_retry_delay = 1  # seconds
 
     def process(self, element: bytes) -> None:
         try:
@@ -34,7 +37,7 @@ def process(self, element: bytes) -> None:
 
             # If we've reached the batch size, send the batch
             if len(self.batch) >= self.batch_size:
-                self._send_batch()
+                self._send_batch_with_retry()
 
         except Exception as e:
             # Log the error but don't fail the pipeline
@@ -43,30 +46,47 @@ def process(self, element: bytes) -> None:
     def finish_bundle(self):
         # Send any remaining messages in the batch
         if self.batch:
-            self._send_batch()
+            self._send_batch_with_retry()
 
-    def _send_batch(self):
-        try:
-            # Convert batch to JSON and compress with gzip
-            json_data = json.dumps(self.batch)
-            compressed_data = gzip.compress(json_data.encode('utf-8'))
-            
-            # Send compressed batch to Better Stack
-            response = requests.post(
-                self.ingesting_url,
-                headers=self.headers,
-                data=compressed_data
-            )
-            
-            if response.status_code != 202:
-                raise Exception(f"Failed to send to Better Stack: {response.text}")
+    def _send_batch_with_retry(self):
+        retry_count = 0
+        retry_delay = self.initial_retry_delay
+        
+        while retry_count < self.max_retries:
+            try:
+                # Convert batch to JSON and compress with gzip
+                json_data = json.dumps(self.batch)
+                compressed_data = gzip.compress(json_data.encode('utf-8'))
 
-            # Clear the batch after successful send
-            self.batch = []
+                # Send compressed batch to Better Stack
+                response = requests.post(
+                    self.ingesting_url,
+                    headers=self.headers,
+                    data=compressed_data
+                )
 
-        except Exception as e:
-            # Log the error but don't fail the pipeline
-            print(f"Error sending batch to Better Stack: {str(e)}")
+                if response.status_code == 202:
+                    # Success - clear the batch and return
+                    self.batch = []
+                    return
+                elif response.status_code == 429:  # Rate limit
+                    retry_after = int(response.headers.get('Retry-After', retry_delay))
+                    print(f"Rate limited. Retrying after {retry_after} seconds...")
+                    time.sleep(retry_after)
+                    retry_count += 1
+                    continue
+                else:
+                    raise Exception(f"Failed to send to Better Stack: {response.text}")
+                    
+            except Exception as e:
+                retry_count += 1
+                if retry_count < self.max_retries:
+                    print(f"Attempt {retry_count} failed: {str(e)}. Retrying in {retry_delay} seconds...")
+                    time.sleep(retry_delay)
+                    retry_delay *= 2  # Exponential backoff
+                else:
+                    print(f"All retry attempts failed. Last error: {str(e)}")
+                    return
 
 def run(argv=None):
     parser = argparse.ArgumentParser()
@@ -97,6 +117,18 @@ def run(argv=None):
         type=int,
         help='Window size in seconds for batching messages'
     )
+    parser.add_argument(
+        '--max_retries',
+        default=3,
+        type=int,
+        help='Maximum number of retry attempts for failed requests'
+    )
+    parser.add_argument(
+        '--initial_retry_delay',
+        default=1,
+        type=int,
+        help='Initial delay between retries in seconds'
+    )
     known_args, pipeline_args = parser.parse_known_args(argv)
 
     pipeline_options = PipelineOptions(