onstructor(onCompleteCallback, trrList) {
    this.onCompleteCallback = onCompleteCallback;
    this.trrList = trrList;
    this.aborted = false;
    this.networkUnstable = false;
    this.captivePortal = false;

    this.domains = [];
    for (let i = 0; i < lazy.kRepeats; ++i) {
      // false-y domain will cause DNSLookup to generate a random one.
      this.domains.push(null);
    }
    this.domains.push(...lazy.kPopularDomains);
    this.totalLookups = this.trrList.length * this.domains.length;
    this.completedLookups = 0;
    this.results = [];
  }

  run() {
    if (this._ran || this._aborted) {
      console.error("Trying to re-run a LookupAggregator.");
      return;
    }

    this._ran = true;
    for (let trr of this.trrList) {
      for (let domain of this.domains) {
        new DNSLookup(
          domain,
          trr,
          (request, record, status, usedDomain, retryCount) => {
            this.results.push({
              domain: usedDomain,
              trr,
              status,
              time: record
                ? record.QueryInterface(Ci.nsIDNSAddrRecord)
                    .trrFetchDurationNetworkOnly
                : -1,
              retryCount,
            });

            this.completedLookups++;
            if (this.completedLookups == this.totalLookups) {
              this.recordResults();
            }
          }
        ).doLookup();
      }
    }
  }

  abort() {
    this.aborted = true;
  }

  markUnstableNetwork() {
    this.networkUnstable = true;
  }

  markCaptivePortal() {
    this.captivePortal = true;
  }

  recordResults() {
    if (this.aborted) {
      return;
    }

    for (let { domain, trr, status, time, retryCount } of this.results) {
      if (
        !(
          lazy.kPopularDomains.includes(domain) ||
          domain.includes(lazy.kCanonicalDomain)
        )
      ) {
        console.error("Expected known domain for reporting, got ", domain);
        return;
      }

      Glean.securityDohTrrPerformance.resolvedRecord.record({
        value: "success",
        domain,
        trr,
        status,
        time,
        retryCount,
        networkUnstable: this.networkUnstable,
        captivePortal: this.captivePortal,
      });
    }

    this.onCompleteCallback();
  }
}

// This class monitors the network and spawns a new LookupAggregator when ready.
// When the network goes down, an ongoing aggregator is aborted and a new one
// spawned next time we get a link, up to 5 times. On the fifth time, we just
// let the aggegator complete and mark it as tainted.
/**
 *
 */
export class TRRRacer {
  constructor(onCompleteCallback, trrList) {
    this._aggregator = null;
    this._retryCount = 0;
    this._complete = false;
    this._onCompleteCallback = onCompleteCallback;
    this._trrList = trrList;
  }

  run() {
    if (
      lazy.gNetworkLinkService.isLinkUp &&
      lazy.gCaptivePortalService.state !=
        lazy.gCaptivePortalService.LOCKED_PORTAL
    ) {
      this._runNewAggregator();
      if (
        lazy.gCaptivePortalService.state ==
        lazy.gCaptivePortalService.UNLOCKED_PORTAL
      ) {
        this._aggregator.markCaptivePortal();
      }
    }

    Services.obs.addObserver(this, "ipc:network:captive-portal-set-state");
    Services.obs.addObserver(this, "network:link-status-changed");
  }

  onComplete() {
    Services.obs.removeObserver(this, "ipc:network:captive-portal-set-state");
    Services.obs.removeObserver(this, "network:link-status-changed");

    this._complete = true;

    if (this._onCompleteCallback) {
      this._onCompleteCallback();
    }
  }

  getFastestTRR(returnRandomDefault = false) {
    if (!this._complete) {
      throw new Error("getFastestTRR: Measurement still running.");
    }

    return this._getFastestTRRFromResults(
      this._aggregator.results,
      returnRandomDefault
    );
  }

  /**
   * Given an array of { trr, time }, returns the trr with smallest mean time.
   * Separate from _getFastestTRR for easy unit-testing.
   *
   * @param {{trr: string[], time: number}[]} results
   * @param {boolean} returnRandomDefault
   * @returns {string}
   *   The TRR with the fastest average time.
   *   If returnRandomDefault is false-y, returns undefined if no valid
   *   times were present in the results. Otherwise, returns one of the
   *   present TRRs at random.
   */
  _getFastestTRRFromResults(results, returnRandomDefault = false) {
    // First, organize the results into a map of TRR -> array of times
    let TRRTimingMap = new Map();
    let TRRErrorCount = new Map();
    for (let { trr, time } of results) {
      if (!TRRTimingMap.has(trr)) {
        TRRTimingMap.set(trr, []);
      }
      if (time != -1) {
        TRRTimingMap.get(trr).push(time);
      } else {
        TRRErrorCount.set(trr, 1 + (TRRErrorCount.get(trr) || 0));
      }
    }

    // Loop through each TRR's array of times, compute the geometric means,
    // and remember the fastest TRR. Geometric mean is a bit more forgiving
    // in the presence of noise (anomalously high values).
    // We don't need the full geometric mean, we simply calculate the arithmetic
    // means in log-space and then compare those values.
    let fastestTRR;
    let fastestAverageTime = -1;
    let trrs = [...TRRTimingMap.keys()];
    for (let trr of trrs) {
      let times = TRRTimingMap.get(trr);
      if (!times.length) {
        continue;
      }

      // Skip TRRs that had an error rate of more than 30%.
      let errorCount = TRRErrorCount.get(trr) || 0;
      let totalResults = times.length + errorCount;
      if (errorCount / totalResults > 0.3) {
        continue;
      }

      // Arithmetic mean in log space. Take log of (a + 1) to ensure we never
      // take log(0) which would be -Infinity.
      let averageTime =
        times.map(a => Math.log(a + 1)).reduce((a, b) => a + b) / times.length;
      if (fastestAverageTime == -1 || averageTime < fastestAverageTime) {
        fastestAverageTime = averageTime;
        fastestTRR = trr;
      }
    }

    if (returnRandomDefault && !fastestTRR) {
      fastestTRR = trrs[Math.floor(Math.random() * trrs.length)];
    }

    return fastestTRR;
  }

  _runNewAggregator() {
    this._aggregator = new LookupAggregator(
      () => this.onComplete(),
      this._trrList
    );
    this._aggregator.run();
    this._retryCount++;
  }

  // When the link goes *down*, or when we detect a locked captive portal, we
  // abort any ongoing LookupAggregator run. When the link goes *up*, or we
  // detect a newly unlocked portal, we start a run if one isn't ongoing.
  observe(subject, topic, data) {
    switch (topic) {
      case "network:link-status-changed":
        if (this._aggregator && data == "down") {
          if (this._retryCount < 5) {
            this._aggregator.abort();
          } else {
            this._aggregator.markUnstableNetwork();
          }
        } else if (
          data == "up" &&
          (!this._aggregator || this._aggregator.aborted)
        ) {
          this._runNewAggregator();
        }
        break;
      case "ipc:network:captive-portal-set-state":
        if (
          this._aggregator &&
          lazy.gCaptivePortalService.state ==
            lazy.gCaptivePortalService.LOCKED_PORTAL
        ) {
          if (this._retryCount < 5) {
            this._aggregator.abort();
          } else {
            this._aggregator.markCaptivePortal();
          }
        } else if (
          lazy.gCaptivePortalService.state ==
            lazy.gCaptivePortalService.UNLOCKED_PORTAL &&
          (!this._aggregator || this._aggregator.aborted)
        ) {
          this._runNewAggregator();
        }
        break;
    }
  }
}
PK