|
| 1 | +'use strict' |
| 2 | + |
| 3 | +const BaseWriter = require('../../../exporters/common/writer') |
| 4 | + |
| 5 | +const FINAL_FLUSH_TIMEOUT_CODE = 'ERR_DD_TEST_OPTIMIZATION_FLUSH_TIMEOUT' |
| 6 | + |
| 7 | +class TestOptimizationRequestTracker { |
| 8 | + #writer |
| 9 | + #pendingRequests = new Set() |
| 10 | + #finalFlushes = new Set() |
| 11 | + #activeFinalFlush |
| 12 | + |
| 13 | + /** |
| 14 | + * Creates request tracking for a Test Optimization writer. |
| 15 | + * |
| 16 | + * @param {BaseWriter} writer |
| 17 | + */ |
| 18 | + constructor (writer) { |
| 19 | + this.#writer = writer |
| 20 | + } |
| 21 | + |
| 22 | + /** |
| 23 | + * Flushes queued payloads and, for a final flush, waits for requests that were |
| 24 | + * already in flight. The absolute deadline prevents Test Optimization from |
| 25 | + * keeping the test process alive indefinitely. |
| 26 | + * |
| 27 | + * @param {(error?: Error) => void} [done] |
| 28 | + * @param {{ deadline?: number }} [options] |
| 29 | + * @returns {void} |
| 30 | + */ |
| 31 | + flush (done, options) { |
| 32 | + if (options?.deadline === undefined) { |
| 33 | + BaseWriter.prototype.flush.call(this.#writer, done, options) |
| 34 | + return |
| 35 | + } |
| 36 | + |
| 37 | + const finalFlush = { |
| 38 | + deadline: options.deadline, |
| 39 | + done: done || (() => {}), |
| 40 | + error: undefined, |
| 41 | + requests: new Set(), |
| 42 | + timeoutId: undefined, |
| 43 | + writerDone: false, |
| 44 | + } |
| 45 | + this.#finalFlushes.add(finalFlush) |
| 46 | + |
| 47 | + for (const pendingRequest of this.#pendingRequests) { |
| 48 | + this.#attachRequest(finalFlush, pendingRequest) |
| 49 | + } |
| 50 | + |
| 51 | + const remaining = Math.max(0, options.deadline - Date.now()) |
| 52 | + finalFlush.timeoutId = setTimeout(() => { |
| 53 | + const error = new Error('Timed out flushing Test Optimization data') |
| 54 | + error.code = FINAL_FLUSH_TIMEOUT_CODE |
| 55 | + |
| 56 | + finalFlush.error ||= error |
| 57 | + finalFlush.writerDone = true |
| 58 | + |
| 59 | + for (const pendingRequest of finalFlush.requests) { |
| 60 | + pendingRequest.finalFlushes.delete(finalFlush) |
| 61 | + if (pendingRequest.finalFlushes.size === 0) { |
| 62 | + pendingRequest.controller.abort(error) |
| 63 | + this.#pendingRequests.delete(pendingRequest) |
| 64 | + } else { |
| 65 | + this.#updateRequestDeadline(pendingRequest) |
| 66 | + } |
| 67 | + } |
| 68 | + finalFlush.requests.clear() |
| 69 | + this.#finishFinalFlush(finalFlush) |
| 70 | + }, remaining) |
| 71 | + |
| 72 | + const previousFinalFlush = this.#activeFinalFlush |
| 73 | + this.#activeFinalFlush = finalFlush |
| 74 | + try { |
| 75 | + BaseWriter.prototype.flush.call(this.#writer, (error) => { |
| 76 | + finalFlush.error ||= error |
| 77 | + finalFlush.writerDone = true |
| 78 | + this.#finishFinalFlush(finalFlush) |
| 79 | + }, options) |
| 80 | + } finally { |
| 81 | + this.#activeFinalFlush = previousFinalFlush |
| 82 | + } |
| 83 | + this.#finishFinalFlush(finalFlush) |
| 84 | + } |
| 85 | + |
| 86 | + /** |
| 87 | + * Sends and tracks a request so a later final flush can wait for or abort it. |
| 88 | + * |
| 89 | + * @param {Function} request |
| 90 | + * @param {Buffer|string|object} data |
| 91 | + * @param {object} options |
| 92 | + * @param {(error: Error|null, result?: string|null, statusCode?: number, |
| 93 | + * headers?: import('node:http').IncomingHttpHeaders) => void} callback |
| 94 | + * @returns {void} |
| 95 | + */ |
| 96 | + send (request, data, options, callback) { |
| 97 | + const controller = new AbortController() |
| 98 | + const requestOptions = { ...options, signal: controller.signal } |
| 99 | + const pendingRequest = { controller, finalFlushes: new Set(), options: requestOptions } |
| 100 | + this.#pendingRequests.add(pendingRequest) |
| 101 | + if (this.#activeFinalFlush) this.#attachRequest(this.#activeFinalFlush, pendingRequest) |
| 102 | + |
| 103 | + request(data, requestOptions, (error, result, statusCode, headers) => { |
| 104 | + if (error) { |
| 105 | + for (const finalFlush of pendingRequest.finalFlushes) finalFlush.error ||= error |
| 106 | + } |
| 107 | + |
| 108 | + try { |
| 109 | + callback(error, result, statusCode, headers) |
| 110 | + } finally { |
| 111 | + this.#pendingRequests.delete(pendingRequest) |
| 112 | + for (const finalFlush of pendingRequest.finalFlushes) { |
| 113 | + finalFlush.requests.delete(pendingRequest) |
| 114 | + this.#finishFinalFlush(finalFlush) |
| 115 | + } |
| 116 | + pendingRequest.finalFlushes.clear() |
| 117 | + } |
| 118 | + }) |
| 119 | + } |
| 120 | + |
| 121 | + /** |
| 122 | + * Associates a request with the final flush boundary that must wait for it. |
| 123 | + * |
| 124 | + * @param {object} finalFlush |
| 125 | + * @param {object} pendingRequest |
| 126 | + * @returns {void} |
| 127 | + */ |
| 128 | + #attachRequest (finalFlush, pendingRequest) { |
| 129 | + finalFlush.requests.add(pendingRequest) |
| 130 | + pendingRequest.finalFlushes.add(finalFlush) |
| 131 | + this.#updateRequestDeadline(pendingRequest) |
| 132 | + } |
| 133 | + |
| 134 | + /** |
| 135 | + * Gives a shared request the latest deadline of the flushes waiting for it. |
| 136 | + * |
| 137 | + * @param {object} pendingRequest |
| 138 | + * @returns {void} |
| 139 | + */ |
| 140 | + #updateRequestDeadline (pendingRequest) { |
| 141 | + let deadline = 0 |
| 142 | + for (const finalFlush of pendingRequest.finalFlushes) { |
| 143 | + deadline = Math.max(deadline, finalFlush.deadline) |
| 144 | + } |
| 145 | + pendingRequest.options.deadline = deadline |
| 146 | + } |
| 147 | + |
| 148 | + /** |
| 149 | + * Completes a final flush callback once its writer and associated requests |
| 150 | + * have settled. |
| 151 | + * |
| 152 | + * @param {object} finalFlush |
| 153 | + * @returns {void} |
| 154 | + */ |
| 155 | + #finishFinalFlush (finalFlush) { |
| 156 | + if (!this.#finalFlushes.has(finalFlush) || !finalFlush.writerDone || finalFlush.requests.size !== 0) return |
| 157 | + |
| 158 | + this.#finalFlushes.delete(finalFlush) |
| 159 | + clearTimeout(finalFlush.timeoutId) |
| 160 | + finalFlush.done(finalFlush.error) |
| 161 | + } |
| 162 | +} |
| 163 | + |
| 164 | +module.exports = TestOptimizationRequestTracker |
0 commit comments