Merge pull request #27 from Firstset/relay-check-improved

Relay check improved

Author: bernatfp

README.md

@@ -174,7 +174,11 @@ The state check subcommand will retrieve all keys registered in CSM, validator c
 
 The relay check subcommand will check if the validator keys are registered with all whitelisted relays in the Lido Relay Allowlist.
 
-`python -m swarm relay-check`
+`python -m swarm relay-check [--detailed] [--pubkey <pubkey>]`
+
+The `--detailed` flag will print a more detailed report with specific relay URLs per validator.
+
+The `--pubkey` flag allows checking a specific validator, showing a detailed report for just that validator.
 
 ### Manual exit
 

swarm/__main__.py

@@ -38,6 +38,8 @@
     parser_exit_monitor.set_defaults(func=exit.automated_exit)
 
     parser_relay_check = subparsers.add_parser('relay-check', help='Check validator relay coverage against allowed mev-boost relays')
+    parser_relay_check.add_argument('--detailed', action='store_true', help='Show detailed registration information')
+    parser_relay_check.add_argument('--pubkey', type=str, help='the public key of the validator to report on')
     parser_relay_check.set_defaults(func=relay_check.check_relays)
 
     args = parser.parse_args()

swarm/relay_check.py

@@ -21,12 +21,13 @@ async def check_validator_registration(session, relay_url, pubkey):
     """Check if a validator is registered with a specific relay"""
     url = f"{relay_url}/relay/v1/data/validator_registration"
     try:
-        async with session.get(url, params={'pubkey': pubkey}) as response:
+        async with session.get(url, params={'pubkey': pubkey}, timeout=5) as response:
             if response.status == 200:
                 return True
             return False
-    except:
-        raise RelayRequestException(f"Failed to request validator registration from {relay_url}")
+    except Exception as e:
+        print(f"Warning: Could not check registration status for {relay_url}: {str(e)}")
+        return None  # Return None to indicate unknown status
 
 async def get_validator_keys_from_csm(config):
     csm = CSM(config)
@@ -39,37 +40,125 @@ async def check_relays(config, args):
 
     # Get whitelisted relays
     relay_tuples = await get_whitelisted_relays(config)
-    # Extract just the URLs from the relay tuples
     relay_urls = [relay[0] for relay in relay_tuples]
-    print(f"Found {len(relay_urls)} whitelisted relays")
+    
+    # Separate mandatory and optional relays
+    mandatory_relays = [relay[0] for relay in relay_tuples if relay[2]]  # relay[2] is the mandatory flag
+    optional_relays = [relay[0] for relay in relay_tuples if not relay[2]]
+    
+    print(f"Found {len(relay_urls)} whitelisted relays ({len(mandatory_relays)} mandatory, {len(optional_relays)} optional)")
 
     # Get validator public keys from CSM
     validator_keys = await get_validator_keys_from_csm(config)
-    print(f"Found {len(validator_keys)} validators in CSM")
+    
+    # If pubkey is specified, only check that validator
+    if hasattr(args, 'pubkey') and args.pubkey:
+        if args.pubkey not in validator_keys:
+            print(f"Error: Validator {args.pubkey} not found in CSM")
+            return
+        validator_keys = [args.pubkey]
+        print(f"Checking specific validator: {args.pubkey}")
+    else:
+        print(f"Found {len(validator_keys)} validators in CSM")
 
     # Check registration for each validator with each relay
     async with aiohttp.ClientSession() as session:
         results = {}
-        for validator in validator_keys:
-            results[validator] = {'total': 0, 'relays': []}
+        
+        # Process validators in chunks to limit concurrency
+        chunk_size = 10
+        for i in range(0, len(validator_keys), chunk_size):
+            validator_chunk = validator_keys[i:i + chunk_size]
+            
+            tasks = []
+            for validator in validator_chunk:
+                results[validator] = {
+                    'mandatory_total': 0,
+                    'optional_total': 0,
+                    'mandatory_relays': [],
+                    'optional_relays': [],
+                    'unknown': []
+                }
+                for relay_url in relay_urls:
+                    tasks.append(check_validator_registration(session, relay_url, validator))
+            
+            chunk_results = await asyncio.gather(*tasks)
+            
+            for validator_idx, validator in enumerate(validator_chunk):
+                for relay_idx, is_registered in enumerate(chunk_results[validator_idx * len(relay_urls):(validator_idx + 1) * len(relay_urls)]):
+                    relay_url = relay_urls[relay_idx]
+                    if is_registered is True:
+                        if relay_url in mandatory_relays:
+                            results[validator]['mandatory_total'] += 1
+                            results[validator]['mandatory_relays'].append(relay_url)
+                        else:
+                            results[validator]['optional_total'] += 1
+                            results[validator]['optional_relays'].append(relay_url)
+                    elif is_registered is None:
+                        results[validator]['unknown'].append(relay_url)
+            
+            if not args.pubkey:  # Only show progress for bulk checks
+                print(f"Processed {min(i + chunk_size, len(validator_keys))}/{len(validator_keys)} validators...")
+        
+        # Print results based on mode
+        if args.pubkey or (hasattr(args, 'detailed') and args.detailed):
+            print_detailed_report(results, relay_tuples, mandatory_relays, optional_relays)
+        else:
+            print_summary_report(results, mandatory_relays, optional_relays)
+
+def print_summary_report(results, mandatory_relays, optional_relays):
+    """Print condensed summary of validator relay registration status"""
+    print("\nRelay Coverage Summary:")
+    print("-" * 50)
+    
+    ok_count = 0
+    not_ok_count = 0
+    
+    for validator, data in results.items():
+        is_ok = data['mandatory_total'] >= 2  # At least 2 mandatory relays required
+        status = "OK" if is_ok else "NOT OK"
+        
+        if is_ok:
+            ok_count += 1
+        else:
+            not_ok_count += 1
 
-            for relay_url in relay_urls:
-                is_registered = await check_validator_registration(session, relay_url, validator)
-                if is_registered:
-                    results[validator]['total'] += 1
-                    results[validator]['relays'].append(relay_url)
+        mandatory_coverage = f"{data['mandatory_total']}/{len(mandatory_relays)}"
+        optional_coverage = f"{data['optional_total']}/{len(optional_relays)}"
+        
+        print(f"Validator {validator} : {status:6} (Mandatory: {mandatory_coverage}, Optional: {optional_coverage})")
+    
+    print(f"\nSummary: {ok_count} validators OK, {not_ok_count} validators NOT OK")
+
+def print_detailed_report(results, relay_tuples, mandatory_relays, optional_relays):
+    """Print detailed report of validator relay registration status"""
+    print("\nDetailed Relay Coverage Report:")
+    print("-" * 50)
+    
+    for validator, data in results.items():
+        is_ok = data['mandatory_total'] >= 2
+        status = "OK" if is_ok else "NOT OK"
+        
+        print(f"\nValidator {validator} : {status}")
+        print(f"Registered with {data['mandatory_total']}/{len(mandatory_relays)} mandatory relays")
+        print(f"Registered with {data['optional_total']}/{len(optional_relays)} optional relays")
+        
+        if data['unknown']:
+            print(f"Could not check status for {len(data['unknown'])} relay(s):")
+            for relay in data['unknown']:
+                relay_info = next(r for r in relay_tuples if r[0] == relay)
+                print(f"  ? {relay} ({relay_info[1]} - {relay_info[3]})")
+        
+        missing_mandatory = set(mandatory_relays) - set(data['mandatory_relays'])
+        if missing_mandatory:
+            print("Missing mandatory registrations:")
+            for relay in missing_mandatory:
+                relay_info = next(r for r in relay_tuples if r[0] == relay)
+                print(f"  - {relay} ({relay_info[1]} - {relay_info[3]})")
 
-        # Print results
-        print("\nRelay coverage report:")
-        print("-" * 50)
-        for validator, data in results.items():
-            coverage_pct = (data['total'] / len(relay_urls)) * 100
-            print(f"\nValidator {validator[:12]}...")
-            print(f"Registered with {data['total']}/{len(relay_urls)} relays ({coverage_pct:.1f}%)")
-            if data['total'] < len(relay_urls):
-                print("Missing registrations for relays:")
-                missing_relays = set(relay_urls) - set(data['relays'])
-                for relay in missing_relays:
-                    # Find the full relay info for prettier printing
-                    relay_info = next(r for r in relay_tuples if r[0] == relay)
-                    print(f"  - {relay} ({relay_info[1]} - {relay_info[3]})")
+        missing_optional = set(optional_relays) - set(data['optional_relays'])
+        if missing_optional:
+            print("Missing optional registrations:")
+            for relay in missing_optional:
+                relay_info = next(r for r in relay_tuples if r[0] == relay)
+                print(f"  - {relay} ({relay_info[1]} - {relay_info[3]})")