Gemini API のリクエスト Hedging で p99 レイテンシのテールを抑える設計

import Foundation
 
actor HedgedGeminiClient {
    private let baseURL = URL(string: "https://generativelanguage.googleapis.com")!
    private let apiKey: String
    private let session: URLSession
 
    // 計測に基づくしきい値 (p95 + α)
    private let hedgeAfterSeconds: TimeInterval = 3.2
    private let totalTimeoutSeconds: TimeInterval = 8.0
 
    init(apiKey: String, session: URLSession = .shared) {
        self.apiKey = apiKey
        self.session = session
    }
 
    func generateContent(prompt: String, model: String = "gemini-2.5-flash") async throws -> String {
        let request = try buildRequest(prompt: prompt, model: model)
 
        return try await withThrowingTaskGroup(of: String.self) { group in
            // 1 本目: 通常リクエスト
            group.addTask { [session] in
                try await Self.send(request: request, session: session)
            }
 
            // 2 本目: hedgeAfterSeconds 待ってから発火
            group.addTask { [hedgeAfterSeconds, session] in
                try await Task.sleep(nanoseconds: UInt64(hedgeAfterSeconds * 1_000_000_000))
                try Task.checkCancellation()
                return try await Self.send(request: request, session: session)
            }
 
            // 全体タイムアウト
            group.addTask { [totalTimeoutSeconds] in
                try await Task.sleep(nanoseconds: UInt64(totalTimeoutSeconds * 1_000_000_000))
                throw GeminiError.timeout
            }
 
            // 先勝ち: 最初の成功レスポンスを採用し残りをキャンセル
            for try await result in group {
                group.cancelAll()
                return result
            }
            throw GeminiError.noResponse
        }
    }
 
    private func buildRequest(prompt: String, model: String) throws -> URLRequest {
        let url = baseURL
            .appendingPathComponent("v1beta/models/\(model):generateContent")
            .appending(queryItems: [URLQueryItem(name: "key", value: apiKey)])
        var req = URLRequest(url: url)
        req.httpMethod = "POST"
        req.setValue("application/json", forHTTPHeaderField: "Content-Type")
        let body: [String: Any] = [
            "contents": [["parts": [["text": prompt]]]]
        ]
        req.httpBody = try JSONSerialization.data(withJSONObject: body)
        // 個別リクエストにも控えめなタイムアウトを設定
        req.timeoutInterval = 7.0
        return req
    }
 
    private static func send(request: URLRequest, session: URLSession) async throws -> String {
        let (data, response) = try await session.data(for: request)
        guard let http = response as? HTTPURLResponse, (200..<300).contains(http.statusCode) else {
            throw GeminiError.invalidResponse
        }
        let json = try JSONSerialization.jsonObject(with: data) as? [String: Any]
        guard let candidates = json?["candidates"] as? [[String: Any]],
              let content = candidates.first?["content"] as? [String: Any],
              let parts = content["parts"] as? [[String: Any]],
              let text = parts.first?["text"] as? String else {
            throw GeminiError.invalidResponse
        }
        return text
    }
 
    enum GeminiError: Error { case timeout, noResponse, invalidResponse }
}

type HedgeOptions = {
  hedgeAfterMs?: number;
  totalTimeoutMs?: number;
  model?: string;
};
 
export async function hedgedGenerate(
  prompt: string,
  apiKey: string,
  opts: HedgeOptions = {}
): Promise<string> {
  const hedgeAfterMs = opts.hedgeAfterMs ?? 3200;
  const totalTimeoutMs = opts.totalTimeoutMs ?? 8000;
  const model = opts.model ?? "gemini-2.5-flash";
 
  const url = `https://generativelanguage.googleapis.com/v1beta/models/${model}:generateContent?key=${apiKey}`;
  const body = JSON.stringify({ contents: [{ parts: [{ text: prompt }] }] });
  const init: RequestInit = {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body,
  };
 
  const controllers = [new AbortController(), new AbortController()];
 
  const send = async (idx: number, delayMs: number): Promise<string> => {
    if (delayMs > 0) await new Promise((r) => setTimeout(r, delayMs));
    const res = await fetch(url, { ...init, signal: controllers[idx].signal });
    if (!res.ok) throw new Error(`HTTP ${res.status}`);
    const json = await res.json<any>();
    return json?.candidates?.[0]?.content?.parts?.[0]?.text ?? "";
  };
 
  const timeout = new Promise<never>((_, reject) =>
    setTimeout(() => reject(new Error("hedged-timeout")), totalTimeoutMs)
  );
 
  try {
    const winner = await Promise.race([
      send(0, 0),
      send(1, hedgeAfterMs),
      timeout,
    ]);
    return winner;
  } finally {
    // 勝者が決まり次第、両方の AbortController を発火する
    controllers.forEach((c) => c.abort());
  }
}

// 広告表示と同時に応答を先行確定させる
async function prefetchDuringAd(prompt: string): Promise<string> {
  const AD_BUDGET_MS = 4800;   // 広告の平均表示時間
  const SAFETY_MS = 1000;      // 閉じる瞬間のばらつきを吸収
  const slo = AD_BUDGET_MS - SAFETY_MS; // 3800ms を上限に応答を取り切る
 
  const result = await hedgedGenerate(prompt, {
    hedgeAfterMs: Math.min(1800, slo / 2), // SLO の半分で 2 本目
    deadlineMs: slo,
  });
  return result; // 広告が閉じた時には result が手元にある
}