1import * as xray from "@distilled.cloud/aws/xray";
2import * as Data from "effect/Data";
3import * as Effect from "effect/Effect";
4import * as Option from "effect/Option";
5import * as Stream from "effect/Stream";
6import { Unowned } from "../../AdoptPolicy.ts";
7import { isResolved } from "../../Diff.ts";
8import { createPhysicalName } from "../../PhysicalName.ts";
9import * as Provider from "../../Provider.ts";
10import { Resource } from "../../Resource.ts";
11import {
12 createInternalTags,
13 createTagsList,
14 diffTags,
15 hasAlchemyTags,
16} from "../../Tags.ts";
17import { AWSEnvironment } from "../Environment.ts";
18import type { Providers } from "../Providers.ts";
19
20export interface SamplingRuleProps {
21 /**
22 * Name of the sampling rule (1-32 characters). `Default` is reserved by
23 * X-Ray for the built-in fallback rule.
24 *
25 * Changing the name replaces the rule.
26 * @default ${app}-${stage}-${id}
27 */
28 ruleName?: string;
29 /**
30 * The priority of the sampling rule. X-Ray evaluates rules in ascending
31 * order of priority (`1` - `9999`) and applies the first rule that matches.
32 */
33 priority: number;
34 /**
35 * The percentage of matching requests to instrument, after the reservoir is
36 * exhausted (`0.0` - `1.0`, e.g. `0.05` for 5%).
37 */
38 fixedRate: number;
39 /**
40 * A fixed number of matching requests to instrument per second, before
41 * applying the fixed rate. The reservoir is not used directly by services,
42 * but applies to all services using the rule collectively.
43 * @default 0
44 */
45 reservoirSize?: number;
46 /**
47 * Matches the `name` that the service uses to identify itself in segments.
48 * @default "*"
49 */
50 serviceName?: string;
51 /**
52 * Matches the `origin` that the service uses to identify its type in
53 * segments (e.g. `AWS::Lambda::Function`).
54 * @default "*"
55 */
56 serviceType?: string;
57 /**
58 * Matches the hostname from a request URL.
59 * @default "*"
60 */
61 host?: string;
62 /**
63 * Matches the HTTP method of a request.
64 * @default "*"
65 */
66 httpMethod?: string;
67 /**
68 * Matches the path from a request URL.
69 * @default "*"
70 */
71 urlPath?: string;
72 /**
73 * Matches the ARN of the Amazon Web Services resource on which the service
74 * runs.
75 * @default "*"
76 */
77 resourceArn?: string;
78 /**
79 * Matches attributes derived from the request (segment annotations).
80 */
81 attributes?: Record<string, string>;
82 /**
83 * Tags to apply to the sampling rule. Merged with internal Alchemy tags.
84 */
85 tags?: Record<string, string>;
86}
87
88export interface SamplingRule extends Resource<
89 "AWS.XRay.SamplingRule",
90 SamplingRuleProps,
91 {
92 /**
93 * Name of the sampling rule.
94 */
95 ruleName: string;
96 /**
97 * ARN of the sampling rule
98 * (`arn:aws:xray:{region}:{account}:sampling-rule/{name}`).
99 */
100 ruleArn: string;
101 },
102 never,
103 Providers
104> {}
105
106/**
107 * An AWS X-Ray sampling rule that controls which requests are recorded as
108 * traces by instrumented applications.
109 *
110 * X-Ray evaluates sampling rules in ascending priority order for each
111 * request. The first matching rule borrows from its reservoir, then applies
112 * the fixed rate.
113 * ### Creating Sampling Rules
114 * **Example:** Sample all requests to a service
115 * ```typescript
116 * import * as XRay from "alchemy/AWS/XRay";
117 *
118 * const rule = yield* XRay.SamplingRule("SampleEverything", {
119 * priority: 100,
120 * fixedRate: 1.0,
121 * reservoirSize: 5,
122 * serviceName: "my-api-*",
123 * });
124 * ```
125 *
126 * **Example:** Low-rate sampling for a noisy endpoint
127 * ```typescript
128 * const rule = yield* XRay.SamplingRule("HealthChecks", {
129 * priority: 10,
130 * fixedRate: 0.01,
131 * urlPath: "/health",
132 * httpMethod: "GET",
133 * });
134 * ```
135 *
136 * **Example:** Match on segment attributes
137 * ```typescript
138 * const rule = yield* XRay.SamplingRule("PremiumTenants", {
139 * priority: 50,
140 * fixedRate: 0.5,
141 * reservoirSize: 1,
142 * attributes: { tier: "premium" },
143 * });
144 * ```
145 *
146 * @resource
147 */
148export const SamplingRule = Resource<SamplingRule>("AWS.XRay.SamplingRule");
149
150/**
151 * Raised when a `SamplingRule` is configured with the reserved rule name
152 * `Default`, which X-Ray uses for the built-in fallback rule.
153 */
154export class XRayReservedRuleName extends Data.TaggedError(
155 "XRayReservedRuleName",
156)<{ message: string }> {}
157
158const validateRuleName = (props: Pick<SamplingRuleProps, "ruleName">) =>
159 props.ruleName === "Default"
160 ? Effect.fail(
161 new XRayReservedRuleName({
162 message:
163 '"Default" is reserved for the built-in X-Ray fallback sampling rule — choose another ruleName.',
164 }),
165 )
166 : Effect.void;
167
168const shallowRecordEqual = (
169 a: Record<string, string | undefined>,
170 b: Record<string, string | undefined>,
171) => {
172 const aKeys = Object.keys(a).filter((k) => a[k] !== undefined);
173 const bKeys = Object.keys(b).filter((k) => b[k] !== undefined);
174 return (
175 aKeys.length === bKeys.length && aKeys.every((key) => a[key] === b[key])
176 );
177};
178
179export const SamplingRuleProvider = () =>
180 Provider.effect(
181 SamplingRule,
182 Effect.gen(function* () {
183 const createRuleName = Effect.fn(function* (
184 id: string,
185 props: Pick<SamplingRuleProps, "ruleName">,
186 ) {
187 // X-Ray rule names are limited to 32 characters.
188 return (
189 props.ruleName ?? (yield* createPhysicalName({ id, maxLength: 32 }))
190 );
191 });
192
193 // The desired wire shape for both create and update, derived purely
194 // from the new props.
195 const desiredRule = (name: string, props: SamplingRuleProps) => ({
196 RuleName: name,
197 ResourceARN: props.resourceArn ?? "*",
198 Priority: props.priority,
199 FixedRate: props.fixedRate,
200 ReservoirSize: props.reservoirSize ?? 0,
201 ServiceName: props.serviceName ?? "*",
202 ServiceType: props.serviceType ?? "*",
203 Host: props.host ?? "*",
204 HTTPMethod: props.httpMethod ?? "*",
205 URLPath: props.urlPath ?? "*",
206 Attributes: props.attributes ?? {},
207 });
208
209 // X-Ray has no GetSamplingRule — observe by enumerating the (small)
210 // rule list and matching on the deterministic rule name.
211 const observeRule = (name: string) =>
212 xray.getSamplingRules.items({}).pipe(
213 Stream.filter((record) => record.SamplingRule?.RuleName === name),
214 Stream.runHead,
215 Effect.map((record) => Option.getOrUndefined(record)?.SamplingRule),
216 );
217
218 const observedTags = (ruleArn: string) =>
219 xray.listTagsForResource.items({ ResourceARN: ruleArn }).pipe(
220 Stream.runCollect,
221 Effect.map((chunk) =>
222 Object.fromEntries(Array.from(chunk).map((t) => [t.Key, t.Value])),
223 ),
224 Effect.catchTag("ResourceNotFoundException", () =>
225 Effect.succeed({} as Record<string, string>),
226 ),
227 );
228
229 return SamplingRule.Provider.of({
230 stables: ["ruleName", "ruleArn"],
231 list: () =>
232 Effect.gen(function* () {
233 const records = yield* xray.getSamplingRules
234 .items({})
235 .pipe(Stream.runCollect);
236 return Array.from(records).flatMap((record) =>
237 record.SamplingRule?.RuleName &&
238 record.SamplingRule.RuleARN &&
239 // The built-in `Default` sampling rule always exists and can
240 // never be deleted — keep it out of enumeration for
241 // account-wide teardown (nuke).
242 record.SamplingRule.RuleName !== "Default"
243 ? [
244 {
245 ruleName: record.SamplingRule.RuleName,
246 ruleArn: record.SamplingRule.RuleARN,
247 },
248 ]
249 : [],
250 );
251 }),
252 read: Effect.fn(function* ({ id, olds, output }) {
253 const ruleName =
254 output?.ruleName ?? (yield* createRuleName(id, olds ?? {}));
255 const found = yield* observeRule(ruleName);
256 if (!found?.RuleARN) return undefined;
257 const attrs = { ruleName, ruleArn: found.RuleARN };
258 const tags = yield* observedTags(found.RuleARN);
259 return (yield* hasAlchemyTags(id, tags)) ? attrs : Unowned(attrs);
260 }),
261 diff: Effect.fn(function* ({ id, news, olds }) {
262 if (!isResolved(news)) return undefined;
263 yield* validateRuleName(news ?? {});
264 const oldName = yield* createRuleName(id, olds ?? {});
265 const newName = yield* createRuleName(id, news ?? {});
266 if (oldName !== newName) {
267 return { action: "replace" } as const;
268 }
269 // fall through: engine default update logic for mutable fields
270 }),
271 reconcile: Effect.fn(function* ({ id, news, output, session }) {
272 yield* validateRuleName(news);
273 const ruleName =
274 output?.ruleName ?? (yield* createRuleName(id, news));
275 const desired = desiredRule(ruleName, news);
276 const internalTags = yield* createInternalTags(id);
277 const desiredTags = { ...news.tags, ...internalTags };
278
279 // 1. OBSERVE — the rule list is authoritative; output is only a
280 // cached identifier.
281 let live = yield* observeRule(ruleName);
282
283 // 2. ENSURE — create when missing; a concurrent create surfaces as
284 // the typed SamplingRuleAlreadyExists tag, which we treat as a
285 // race and re-observe.
286 if (live === undefined) {
287 live = yield* xray
288 .createSamplingRule({
289 SamplingRule: { ...desired, Version: 1 },
290 Tags: createTagsList(desiredTags),
291 })
292 .pipe(
293 Effect.map((r) => r.SamplingRuleRecord?.SamplingRule),
294 Effect.catchTag("SamplingRuleAlreadyExists", () =>
295 observeRule(ruleName),
296 ),
297 );
298 }
299
300 // 3. SYNC — diff every mutable field of the OBSERVED rule against
301 // the desired shape; a single UpdateSamplingRule applies the
302 // full desired state only when something actually drifted.
303 const inSync =
304 live !== undefined &&
305 live.ResourceARN === desired.ResourceARN &&
306 live.Priority === desired.Priority &&
307 live.FixedRate === desired.FixedRate &&
308 live.ReservoirSize === desired.ReservoirSize &&
309 live.ServiceName === desired.ServiceName &&
310 live.ServiceType === desired.ServiceType &&
311 live.Host === desired.Host &&
312 live.HTTPMethod === desired.HTTPMethod &&
313 live.URLPath === desired.URLPath &&
314 shallowRecordEqual(live.Attributes ?? {}, desired.Attributes);
315 if (!inSync) {
316 live = yield* xray
317 .updateSamplingRule({ SamplingRuleUpdate: desired })
318 .pipe(Effect.map((r) => r.SamplingRuleRecord?.SamplingRule));
319 }
320
321 // The ARN shape is deterministic:
322 // arn:aws:xray:{region}:{account}:sampling-rule/{name}
323 const { accountId, region } = yield* AWSEnvironment.current;
324 const ruleArn =
325 live?.RuleARN ??
326 `arn:aws:xray:${region}:${accountId}:sampling-rule/${ruleName}`;
327
328 // 3b. SYNC TAGS — diff against OBSERVED cloud tags so adoption
329 // converges (create-time Tags only apply on first create).
330 const currentTags = yield* observedTags(ruleArn);
331 const { upsert, removed } = diffTags(currentTags, desiredTags);
332 if (upsert.length > 0) {
333 yield* xray.tagResource({ ResourceARN: ruleArn, Tags: upsert });
334 }
335 if (removed.length > 0) {
336 yield* xray.untagResource({
337 ResourceARN: ruleArn,
338 TagKeys: removed,
339 });
340 }
341
342 yield* session.note(ruleName);
343 return { ruleName, ruleArn };
344 }),
345 delete: Effect.fn(function* ({ output }) {
346 yield* xray.deleteSamplingRule({ RuleName: output.ruleName }).pipe(
347 // Typed synthetic tag for a missing rule — idempotent delete.
348 Effect.catchTag("SamplingRuleNotFound", () => Effect.void),
349 );
350 }),
351 });
352 }),
353 );