Setting the file. One moment.
Aggregator · AWS Marketplace Metering · aws/agent-toolkit-for-aws · Skills Docs
Back to the top of the page Query Patterns
70
Creating Amazon Aurora Db Cluster With Instances
93
Routing Traffic With Route53 And CloudFront
Resilience Program Design
Creating API Gateway Stage
Next
Script Cleanup
scripts/ aggregator.py
Python · 199 lines · 9 KB
15 rows asynchronously.
16
17 Idempotency: if the conditional PutItem fails (ConditionalCheckFailedException) the group
18 was already aggregated by a prior visit; skip the put AND skip enqueuing cleanup (inflation
19 guard). collect_group only returns rows still carrying meteringPending, so an already-cleaned
20 group yields no row_keys and no-ops earlier; a genuinely-lost cleanup self-heals on the next
21 scheduled re-visit (rows still pending are re-detected and cleaned), not by re-enqueuing on
22 every duplicate visit.
23
24 Zero-quantity handling: an in-window 0-quantity group is left pending
25 (submitting 0 early would lock the hour at 0 via first-write-wins); at the oldest edge
26 (hour == now-23h) it is aggregated as Quantity 0.
27 """
28
29 import json
30 import os
31 from datetime import datetime, timedelta, timezone
32 from typing import Any, Dict, List
33
34 import boto3
35 from botocore.exceptions import ClientError
36
37 from handlers import metering_core as core
38
39 logger = core.logger
40
41 sqs = boto3.client( "sqs" )
42 CLEANUP_QUEUE_URL = os.environ.get( "CLEANUP_QUEUE_URL" , "" )
43
44
45 def handler (event, context):
46 if core.aggregated_usage_table is None :
47 raise RuntimeError ( "AGGREGATED_USAGE_TABLE not configured" )
48 if not CLEANUP_QUEUE_URL :
49 raise RuntimeError ( "CLEANUP_QUEUE_URL not configured" )
50
51 now = datetime.now(timezone.utc)
52 failures = []
53 for record in event.get( "Records" , []):
54 try :
55 _process_message(json.loads(record[ "body" ]), now)
56 except Exception as e:
57 # Report partial batch failure so only the failed message is retried/DLQ'd.
58 logger.error( f "Aggregation failed for a group: { type (e). __name__ } : { e } " )
59 failures.append({ "itemIdentifier" : record[ "messageId" ]})
60 return { "batchItemFailures" : failures}
61
62
63 def _process_message (msg, now):
64 license_arn = msg[ "licenseArn" ]
65 account_id = msg[ "customerAWSAccountId" ]
66 dimension = msg[ "dimension" ]
67 hour_bucket = msg[ "hourBucket" ]
68
69 group = core.collect_group(license_arn, account_id, dimension, hour_bucket)
70 if not group[ "row_keys" ]:
71 # Nothing pending (already cleaned up by a prior visit) — no-op, idempotent.
72 logger.info(
73 f "No pending rows for group license= { core.mask(license_arn) } "
74 f "dim= { dimension } hour= { hour_bucket } ; skipping"
75 )
76 return
77
78 # In-window zero-quantity group: leave pending (do not aggregate yet). At the oldest
79 # edge it falls through and is aggregated as Quantity 0.
80 oldest_edge = now.replace( minute = 0 , second = 0 , microsecond = 0 ) - timedelta( hours = 23 )
81 if group[ "quantity" ] == 0 and not group[ "reject_reason" ] and group[ "hour_ts" ] > oldest_edge:
82 logger.info(
83 f "In-window zero-quantity group left pending license= { core.mask(license_arn) } "
84 f "dim= { dimension } hour= { hour_bucket } "
85 )
86 return
87
88 reason = core.validate_group(group, now)
89 if reason:
90 _reject(group, reason)
91 return
92
93 written = _write_aggregated(group)
94 if written:
95 # Fresh write — enqueue cleanup carrying the authoritative Aggregated status + sum/count.
96 _enqueue_cleanup(group, stamp_status = True )
97 else :
98 # Duplicate/concurrent visit: a PRIOR aggregator run already wrote the aggregated
99 # record. Do NOT enqueue a (status-less) cleanup here — that only inflates the cleanup
100 # queue. collect_group already filters to rows still carrying meteringPending, so if the
101 # prior cleanup succeeded these rows would have been absent and we'd have no-op'd above;
102 # if a prior cleanup was genuinely lost, the rows still carry meteringPending and are
103 # re-discovered on the NEXT scheduled sweep, which re-detects and cleans them (one-cycle
104 # self-heal) — without a status-less enqueue on every duplicate visit.
105 logger.info(
106 f "Aggregated record already exists (idempotent) for license= { core.mask(license_arn) } "
107 f "dim= { dimension } hour= { hour_bucket } ; NOT re-enqueuing cleanup (inflation guard). "
108 "Any still-pending rows self-heal on the next scheduled re-visit."
109 )
110
111
112 def _write_aggregated (group) -> bool :
113 """Conditional PutItem into aggregated_usage. Returns True if written, False if it
114 already existed (ConditionalCheckFailedException — idempotent commit point)."""
115 item = {
116 "licenseArn" : group[ "license_arn" ],
117 "account_dimension_hour" : core.agg_sort_key(
118 group[ "account_id" ], group[ "dimension" ], group[ "hour_bucket" ]
119 ),
120 "customerAWSAccountId" : group[ "account_id" ],
121 "dimension" : group[ "dimension" ],
122 "hourBucket" : group[ "hour_bucket" ],
123 "quantity" : group[ "quantity" ],
124 "meteringPending" : group[ "hour_bucket" ],
125 }
126 allocations = core.finalize_allocations(group)
127 if allocations:
128 item[ "usageAllocations" ] = allocations
129 if group.get( "customer_identifier" ):
130 item[ "customerIdentifier" ] = group[ "customer_identifier" ]
131 try :
132 core.aggregated_usage_table.put_item(
133 Item = core.audit_item(item),
134 ConditionExpression = "attribute_not_exists(licenseArn)" ,
135 )
136 return True
137 except ClientError as e:
138 if e.response.get( "Error" , {}).get( "Code" ) == "ConditionalCheckFailedException" :
139 return False
140 raise
141
142
143 def _reject (group, reason):
144 """Reject-and-finalize CLIENT-SIDE: terminal status + reason on every raw
145 row of the group, clear meteringPending, WARN, emit the UsageRecordRejected metric.
146 The rejected group is never aggregated or submitted. Writes use the shared bounded
147 thread pool + backoff (a rejected group is up to ~3600 same-licenseArn-partition
148 UpdateItems — same hot-partition guard as cleanup), and raise on any unrecoverable
149 failure so the message is retried (SQS redelivery/DLQ), never silently half-rejected."""
150 logger.warning(
151 "Rejecting usage group client-side "
152 f "(license= { core.mask(group[ 'license_arn' ]) } dim= { group[ 'dimension' ] } "
153 f "hour= { group[ 'hour_bucket' ] } ): { reason } "
154 )
155 core.emf_metric( "UsageRecordRejected" , 1 , reason = reason)
156 core.update_rows_bounded(
157 group[ "license_arn" ],
158 group[ "row_keys" ],
159 "REMOVE meteringPending SET meteringStatus = :st, meteringStatusReason = :rsn" ,
160 { ":st" : core. STATUS_REJECTED , ":rsn" : reason},
161 )
162
163
164 def _enqueue_cleanup (group, stamp_status):
165 """Enqueue cleanup messages carrying ≤100 raw row keys each (PK licenseArn + SKs).
166
167 stamp_status=True (rule a — THIS invocation freshly wrote the aggregated record):
168 include meteringStatus="Aggregated", totalQuantity, and recordCount so the cleanup
169 Lambda stamps the authoritative "Aggregated total quantity <sum> from <count> raw
170 usage records" reason. Only the writer knows the true sum/count.
171 stamp_status=False (rule b — the record already existed; duplicate/concurrent visit):
172 carry ONLY the row keys and NO status — cleanup clears meteringPending only, never
173 stamping a sum that may differ from what the true writer persisted.
174
175 Uses SendMessageBatch (≤10 entries/call, like the discoverer); a partial Failed result
176 raises so the message is retried (a lost cleanup self-heals on the next hourly re-visit)."""
177 license_arn = group[ "license_arn" ]
178 keys = group[ "row_keys" ]
179 entries: List[Dict[ str , Any]] = []
180 for i in range ( 0 , len (keys), core. CLEANUP_KEYS_PER_MESSAGE ):
181 chunk = keys[i : i + core. CLEANUP_KEYS_PER_MESSAGE ]
182 body = { "licenseArn" : license_arn, "rowKeys" : chunk}
183 if stamp_status:
184 body[ "meteringStatus" ] = "Aggregated"
185 body[ "totalQuantity" ] = group[ "quantity" ]
186 body[ "recordCount" ] = len (keys)
187 entries.append({ "Id" : str ( len (entries)), "MessageBody" : json.dumps(body)})
188 if len (entries) == 10 : # SQS SendMessageBatch max 10 entries/call
189 _flush_cleanup_batch(entries)
190 entries = []
191 if entries:
192 _flush_cleanup_batch(entries)
193
194
195 def _flush_cleanup_batch (entries):
196 resp = sqs.send_message_batch( QueueUrl = CLEANUP_QUEUE_URL , Entries = entries)
197 failed = resp.get( "Failed" , [])
198 if failed:
199 raise RuntimeError ( f "SQS SendMessageBatch had { len (failed) } cleanup failure(s): { failed } " )