|
17 | 17 | import asyncio |
18 | 18 | import secrets |
19 | 19 | import time |
| 20 | +from awslabs.aws_healthomics_mcp_server.consts import ( |
| 21 | + BUFFER_EFFICIENCY_HIGH_THRESHOLD, |
| 22 | + BUFFER_EFFICIENCY_LOW_THRESHOLD, |
| 23 | + COMPLEXITY_MULTIPLIER_ASSOCIATED_FILES, |
| 24 | + COMPLEXITY_MULTIPLIER_BUFFER_OVERFLOW, |
| 25 | + COMPLEXITY_MULTIPLIER_FILE_TYPE_FILTER, |
| 26 | + COMPLEXITY_MULTIPLIER_HIGH_EFFICIENCY, |
| 27 | + COMPLEXITY_MULTIPLIER_LOW_EFFICIENCY, |
| 28 | + CURSOR_PAGINATION_BUFFER_THRESHOLD, |
| 29 | + CURSOR_PAGINATION_PAGE_THRESHOLD, |
| 30 | + MAX_SEARCH_RESULTS_LIMIT, |
| 31 | + S3_CACHE_CLEANUP_PROBABILITY, |
| 32 | +) |
20 | 33 | from awslabs.aws_healthomics_mcp_server.models import ( |
21 | 34 | GenomicsFile, |
22 | 35 | GenomicsFileResult, |
@@ -356,8 +369,10 @@ async def search_paginated( |
356 | 369 | ) |
357 | 370 | self._cache_pagination_state(cache_key, cache_entry) |
358 | 371 |
|
359 | | - # Clean up expired cache entries periodically |
360 | | - if secrets.randbelow(20) == 0: # 5% chance to clean up cache |
| 372 | + # Clean up expired cache entries periodically (reduced frequency due to size-based cleanup) |
| 373 | + if ( |
| 374 | + secrets.randbelow(100) == 0 |
| 375 | + ): # Probability defined by PAGINATION_CACHE_CLEANUP_PROBABILITY |
361 | 376 | try: |
362 | 377 | self.cleanup_expired_pagination_cache() |
363 | 378 | except Exception as e: |
@@ -424,8 +439,8 @@ def _validate_search_request(self, request: GenomicsFileSearchRequest) -> None: |
424 | 439 | if request.max_results <= 0: |
425 | 440 | raise ValueError('max_results must be greater than 0') |
426 | 441 |
|
427 | | - if request.max_results > 10000: |
428 | | - raise ValueError('max_results cannot exceed 10000') |
| 442 | + if request.max_results > MAX_SEARCH_RESULTS_LIMIT: |
| 443 | + raise ValueError(f'max_results cannot exceed {MAX_SEARCH_RESULTS_LIMIT}') |
429 | 444 |
|
430 | 445 | # Validate file_type if provided |
431 | 446 | if request.file_type: |
@@ -489,10 +504,11 @@ async def _execute_parallel_searches( |
489 | 504 | else: |
490 | 505 | logger.warning(f'Unexpected result type from {storage_system}: {type(result)}') |
491 | 506 |
|
492 | | - # Periodically clean up expired cache entries (approximately every 10th search) |
| 507 | + # Periodically clean up expired cache entries (reduced frequency due to size-based cleanup) |
493 | 508 | if ( |
494 | | - secrets.randbelow(10) == 0 and self.s3_engine is not None |
495 | | - ): # 10% chance to clean up cache |
| 509 | + secrets.randbelow(100 // S3_CACHE_CLEANUP_PROBABILITY) == 0 |
| 510 | + and self.s3_engine is not None |
| 511 | + ): # Probability defined by S3_CACHE_CLEANUP_PROBABILITY |
496 | 512 | try: |
497 | 513 | self.s3_engine.cleanup_expired_cache_entries() |
498 | 514 | except Exception as e: |
@@ -1003,6 +1019,10 @@ def _cache_pagination_state(self, cache_key: str, entry: 'PaginationCacheEntry') |
1003 | 1019 | if not hasattr(self, '_pagination_cache'): |
1004 | 1020 | self._pagination_cache = {} |
1005 | 1021 |
|
| 1022 | + # Check if we need to clean up before adding |
| 1023 | + if len(self._pagination_cache) >= self.config.max_pagination_cache_size: |
| 1024 | + self._cleanup_pagination_cache_by_size() |
| 1025 | + |
1006 | 1026 | entry.update_timestamp() |
1007 | 1027 | self._pagination_cache[cache_key] = entry |
1008 | 1028 | logger.debug(f'Cached pagination state for key: {cache_key}') |
@@ -1030,26 +1050,26 @@ def _optimize_buffer_size( |
1030 | 1050 |
|
1031 | 1051 | # File type filtering reduces complexity |
1032 | 1052 | if request.file_type: |
1033 | | - complexity_multiplier *= 0.8 |
| 1053 | + complexity_multiplier *= COMPLEXITY_MULTIPLIER_FILE_TYPE_FILTER |
1034 | 1054 |
|
1035 | 1055 | # Associated files increase complexity |
1036 | 1056 | if request.include_associated_files: |
1037 | | - complexity_multiplier *= 1.2 |
| 1057 | + complexity_multiplier *= COMPLEXITY_MULTIPLIER_ASSOCIATED_FILES |
1038 | 1058 |
|
1039 | 1059 | # Adjust based on historical metrics |
1040 | 1060 | if metrics: |
1041 | 1061 | # If we had buffer overflows, increase buffer size |
1042 | 1062 | if metrics.buffer_overflows > 0: |
1043 | | - complexity_multiplier *= 1.5 |
| 1063 | + complexity_multiplier *= COMPLEXITY_MULTIPLIER_BUFFER_OVERFLOW |
1044 | 1064 |
|
1045 | 1065 | # If efficiency was low, increase buffer size |
1046 | 1066 | efficiency_ratio = metrics.total_results_fetched / max( |
1047 | 1067 | metrics.total_objects_scanned, 1 |
1048 | 1068 | ) |
1049 | | - if efficiency_ratio < 0.1: # Less than 10% efficiency |
1050 | | - complexity_multiplier *= 2.0 |
1051 | | - elif efficiency_ratio > 0.5: # More than 50% efficiency |
1052 | | - complexity_multiplier *= 0.8 |
| 1069 | + if efficiency_ratio < BUFFER_EFFICIENCY_LOW_THRESHOLD: |
| 1070 | + complexity_multiplier *= COMPLEXITY_MULTIPLIER_LOW_EFFICIENCY |
| 1071 | + elif efficiency_ratio > BUFFER_EFFICIENCY_HIGH_THRESHOLD: |
| 1072 | + complexity_multiplier *= COMPLEXITY_MULTIPLIER_HIGH_EFFICIENCY |
1053 | 1073 |
|
1054 | 1074 | optimized_size = int(base_buffer_size * complexity_multiplier) |
1055 | 1075 |
|
@@ -1098,9 +1118,63 @@ def _should_use_cursor_pagination( |
1098 | 1118 | """ |
1099 | 1119 | # Use cursor pagination for large buffer sizes or high page numbers |
1100 | 1120 | return self.config.enable_cursor_based_pagination and ( |
1101 | | - request.pagination_buffer_size > 5000 or global_token.page_number > 10 |
| 1121 | + request.pagination_buffer_size > CURSOR_PAGINATION_BUFFER_THRESHOLD |
| 1122 | + or global_token.page_number > CURSOR_PAGINATION_PAGE_THRESHOLD |
| 1123 | + ) |
| 1124 | + |
| 1125 | + def _cleanup_pagination_cache_by_size(self) -> None: |
| 1126 | + """Clean up pagination cache when it exceeds max size, prioritizing expired entries first. |
| 1127 | +
|
| 1128 | + Strategy: |
| 1129 | + 1. First: Remove all expired entries (regardless of age) |
| 1130 | + 2. Then: If still over size limit, remove oldest non-expired entries |
| 1131 | + """ |
| 1132 | + if not hasattr(self, '_pagination_cache'): |
| 1133 | + return |
| 1134 | + |
| 1135 | + if len(self._pagination_cache) < self.config.max_pagination_cache_size: |
| 1136 | + return |
| 1137 | + |
| 1138 | + target_size = int( |
| 1139 | + self.config.max_pagination_cache_size * self.config.cache_cleanup_keep_ratio |
1102 | 1140 | ) |
1103 | 1141 |
|
| 1142 | + # Separate expired and valid entries |
| 1143 | + expired_items = [] |
| 1144 | + valid_items = [] |
| 1145 | + |
| 1146 | + for key, entry in self._pagination_cache.items(): |
| 1147 | + if entry.is_expired(self.config.pagination_cache_ttl_seconds): |
| 1148 | + expired_items.append((key, entry)) |
| 1149 | + else: |
| 1150 | + valid_items.append((key, entry)) |
| 1151 | + |
| 1152 | + # Phase 1: Remove all expired items first |
| 1153 | + expired_count = len(expired_items) |
| 1154 | + for key, _ in expired_items: |
| 1155 | + del self._pagination_cache[key] |
| 1156 | + |
| 1157 | + # Phase 2: If still over target size, remove oldest valid items |
| 1158 | + remaining_count = len(self._pagination_cache) |
| 1159 | + additional_removals = 0 |
| 1160 | + |
| 1161 | + if remaining_count > target_size: |
| 1162 | + # Sort valid items by timestamp (oldest first) |
| 1163 | + valid_items.sort(key=lambda x: x[1].timestamp) |
| 1164 | + additional_to_remove = remaining_count - target_size |
| 1165 | + |
| 1166 | + for i in range(min(additional_to_remove, len(valid_items))): |
| 1167 | + key, _ = valid_items[i] |
| 1168 | + if key in self._pagination_cache: # Double-check key still exists |
| 1169 | + del self._pagination_cache[key] |
| 1170 | + additional_removals += 1 |
| 1171 | + |
| 1172 | + total_removed = expired_count + additional_removals |
| 1173 | + if total_removed > 0: |
| 1174 | + logger.debug( |
| 1175 | + f'Smart pagination cache cleanup: removed {expired_count} expired + {additional_removals} oldest valid = {total_removed} total entries, {len(self._pagination_cache)} remaining' |
| 1176 | + ) |
| 1177 | + |
1104 | 1178 | def cleanup_expired_pagination_cache(self) -> None: |
1105 | 1179 | """Clean up expired pagination cache entries to prevent memory leaks.""" |
1106 | 1180 | if not hasattr(self, '_pagination_cache'): |
@@ -1136,10 +1210,14 @@ def get_pagination_cache_stats(self) -> Dict[str, Any]: |
1136 | 1210 | 'total_entries': len(self._pagination_cache), |
1137 | 1211 | 'valid_entries': valid_entries, |
1138 | 1212 | 'ttl_seconds': self.config.pagination_cache_ttl_seconds, |
| 1213 | + 'max_cache_size': self.config.max_pagination_cache_size, |
| 1214 | + 'cache_utilization': len(self._pagination_cache) |
| 1215 | + / self.config.max_pagination_cache_size, |
1139 | 1216 | 'config': { |
1140 | 1217 | 'enable_cursor_pagination': self.config.enable_cursor_based_pagination, |
1141 | 1218 | 'max_buffer_size': self.config.max_pagination_buffer_size, |
1142 | 1219 | 'min_buffer_size': self.config.min_pagination_buffer_size, |
1143 | 1220 | 'enable_metrics': self.config.enable_pagination_metrics, |
| 1221 | + 'cache_cleanup_keep_ratio': self.config.cache_cleanup_keep_ratio, |
1144 | 1222 | }, |
1145 | 1223 | } |
0 commit comments