feat: improve logging on token refresh logic (#410)

Guilherme Souza · web-flow · commit a8ed053c96ea · 2024-06-05T00:56:08.000+09:00
diff --git a/Sources/Auth/Internal/SessionManager.swift b/Sources/Auth/Internal/SessionManager.swift
@@ -32,57 +32,62 @@ private actor LiveSessionManager {
   private var scheduledNextRefreshTask: Task<Void, Never>?
 
   func session() async throws -> Session {
-    guard let currentSession = try storage.getSession() else {
-      throw AuthError.sessionNotFound
-    }
+    try await trace(using: logger) {
+      guard let currentSession = try storage.getSession() else {
+        throw AuthError.sessionNotFound
+      }
 
-    if !currentSession.isExpired {
-      scheduleNextTokenRefresh(currentSession)
+      if !currentSession.isExpired {
+        await scheduleNextTokenRefresh(currentSession)
 
-      return currentSession
-    }
+        return currentSession
+      }
 
-    return try await refreshSession(currentSession.refreshToken)
+      return try await refreshSession(currentSession.refreshToken)
+    }
   }
 
   func refreshSession(_ refreshToken: String) async throws -> Session {
-    logger?.debug("begin")
-    defer { logger?.debug("end") }
-
-    if let inFlightRefreshTask {
-      logger?.debug("refresh already in flight")
-      return try await inFlightRefreshTask.value
-    }
-
-    inFlightRefreshTask = Task {
-      logger?.debug("refresh task started")
-
-      defer {
-        inFlightRefreshTask = nil
-        logger?.debug("refresh task ended")
+    try await SupabaseLoggerTaskLocal.$additionalContext.withValue(
+      merging: ["refreshID": .string(UUID().uuidString)]
+    ) {
+      try await trace(using: logger) {
+        if let inFlightRefreshTask {
+          logger?.debug("refresh already in flight")
+          return try await inFlightRefreshTask.value
+        }
+
+        inFlightRefreshTask = Task {
+          logger?.debug("refresh task started")
+
+          defer {
+            inFlightRefreshTask = nil
+            logger?.debug("refresh task ended")
+          }
+
+          let session = try await api.execute(
+            HTTPRequest(
+              url: configuration.url.appendingPathComponent("token"),
+              method: .post,
+              query: [
+                URLQueryItem(name: "grant_type", value: "refresh_token"),
+              ],
+              body: configuration.encoder.encode(UserCredentials(refreshToken: refreshToken))
+            )
+          )
+          .decoded(as: Session.self, decoder: configuration.decoder)
+
+          update(session)
+          eventEmitter.emit(.tokenRefreshed, session: session)
+
+          await scheduleNextTokenRefresh(session)
+
+          return session
+        }
+
+        return try await inFlightRefreshTask!.value
       }
-
-      let session = try await api.execute(
-        HTTPRequest(
-          url: configuration.url.appendingPathComponent("token"),
-          method: .post,
-          query: [
-            URLQueryItem(name: "grant_type", value: "refresh_token"),
-          ],
-          body: configuration.encoder.encode(UserCredentials(refreshToken: refreshToken))
-        )
-      )
-      .decoded(as: Session.self, decoder: configuration.decoder)
-
-      update(session)
-      eventEmitter.emit(.tokenRefreshed, session: session)
-
-      scheduleNextTokenRefresh(session)
-
-      return session
     }
-
-    return try await inFlightRefreshTask!.value
   }
 
   func update(_ session: Session) {
@@ -101,37 +106,41 @@ private actor LiveSessionManager {
     }
   }
 
-  private func scheduleNextTokenRefresh(_ refreshedSession: Session, source: StaticString = #function) {
-    logger?.debug("source: \(source)")
+  private func scheduleNextTokenRefresh(_ refreshedSession: Session, caller: StaticString = #function) async {
+    await SupabaseLoggerTaskLocal.$additionalContext.withValue(
+      merging: ["caller": .string("\(caller)")]
+    ) {
+      guard configuration.autoRefreshToken else {
+        logger?.debug("auto refresh token disabled")
+        return
+      }
 
-    guard configuration.autoRefreshToken else {
-      logger?.debug("auto refresh token disabled")
-      return
-    }
+      guard scheduledNextRefreshTask == nil else {
+        logger?.debug("refresh task already scheduled")
+        return
+      }
 
-    guard scheduledNextRefreshTask == nil else {
-      logger?.debug("source: \(source) refresh task already scheduled")
-      return
-    }
+      scheduledNextRefreshTask = Task {
+        await trace(using: logger) {
+          defer { scheduledNextRefreshTask = nil }
 
-    scheduledNextRefreshTask = Task {
-      defer { scheduledNextRefreshTask = nil }
+          let expiresAt = Date(timeIntervalSince1970: refreshedSession.expiresAt)
+          let expiresIn = expiresAt.timeIntervalSinceNow
 
-      let expiresAt = Date(timeIntervalSince1970: refreshedSession.expiresAt)
-      let expiresIn = expiresAt.timeIntervalSinceNow
+          // if expiresIn < 0, it will refresh right away.
+          let timeToRefresh = max(expiresIn * 0.9, 0)
 
-      // if expiresIn < 0, it will refresh right away.
-      let timeToRefresh = max(expiresIn * 0.9, 0)
+          logger?.debug("scheduled next token refresh in: \(timeToRefresh)s")
 
-      logger?.debug("source: \(source) scheduled next token refresh in: \(timeToRefresh)s")
+          try? await Task.sleep(nanoseconds: NSEC_PER_SEC * UInt64(timeToRefresh))
 
-      try? await Task.sleep(nanoseconds: NSEC_PER_SEC * UInt64(timeToRefresh))
+          if Task.isCancelled {
+            return
+          }
 
-      if Task.isCancelled {
-        return
+          _ = try? await refreshSession(refreshedSession.refreshToken)
+        }
       }
-
-      _ = try? await refreshSession(refreshedSession.refreshToken)
     }
   }
 }
diff --git a/Sources/_Helpers/HTTP/LoggerInterceptor.swift b/Sources/_Helpers/HTTP/LoggerInterceptor.swift
@@ -19,28 +19,30 @@ package struct LoggerInterceptor: HTTPClientInterceptor {
     next: @Sendable (HTTPRequest) async throws -> HTTPResponse
   ) async throws -> HTTPResponse {
     let id = UUID().uuidString
-    logger.verbose(
-      """
-      Request [\(id)]: \(request.method.rawValue) \(request.url.absoluteString
-        .removingPercentEncoding ?? "")
-      Body: \(stringfy(request.body))
-      """
-    )
-
-    do {
-      let response = try await next(request)
+    return try await SupabaseLoggerTaskLocal.$additionalContext.withValue(merging: ["requestID": .string(id)]) {
       logger.verbose(
         """
-        Response [\(id)]: Status code: \(response.statusCode) Content-Length: \(
-          response.underlyingResponse.expectedContentLength
-        )
-        Body: \(stringfy(response.data))
+        Request: \(request.method.rawValue) \(request.url.absoluteString
+          .removingPercentEncoding ?? "")
+        Body: \(stringfy(request.body))
         """
       )
-      return response
-    } catch {
-      logger.error("Response [\(id)]: Failure \(error)")
-      throw error
+
+      do {
+        let response = try await next(request)
+        logger.verbose(
+          """
+          Response: Status code: \(response.statusCode) Content-Length: \(
+            response.underlyingResponse.expectedContentLength
+          )
+          Body: \(stringfy(response.data))
+          """
+        )
+        return response
+      } catch {
+        logger.error("Response: Failure \(error)")
+        throw error
+      }
     }
   }
 }
diff --git a/Sources/_Helpers/SupabaseLogger.swift b/Sources/_Helpers/SupabaseLogger.swift
@@ -16,6 +16,13 @@ public enum SupabaseLogLevel: Int, Codable, CustomStringConvertible, Sendable {
   }
 }
 
+@usableFromInline
+package enum SupabaseLoggerTaskLocal {
+  @TaskLocal
+  @usableFromInline
+  package static var additionalContext: JSONObject = [:]
+}
+
 public struct SupabaseLogMessage: Codable, CustomStringConvertible, Sendable {
   public let system: String
   public let level: SupabaseLogLevel
@@ -24,6 +31,7 @@ public struct SupabaseLogMessage: Codable, CustomStringConvertible, Sendable {
   public let function: String
   public let line: UInt
   public let timestamp: TimeInterval
+  public let additionalContext: JSONObject
 
   @usableFromInline
   init(
@@ -33,7 +41,8 @@ public struct SupabaseLogMessage: Codable, CustomStringConvertible, Sendable {
     fileID: String,
     function: String,
     line: UInt,
-    timestamp: TimeInterval
+    timestamp: TimeInterval,
+    additionalContext: JSONObject
   ) {
     self.system = system
     self.level = level
@@ -42,12 +51,17 @@ public struct SupabaseLogMessage: Codable, CustomStringConvertible, Sendable {
     self.function = function
     self.line = line
     self.timestamp = timestamp
+    self.additionalContext = additionalContext
   }
 
   public var description: String {
     let date = iso8601Formatter.string(from: Date(timeIntervalSince1970: timestamp))
     let file = fileID.split(separator: ".", maxSplits: 1).first.map(String.init) ?? fileID
-    return "\(date) [\(level)] [\(system)] [\(file).\(function):\(line)] \(message)"
+    var description = "\(date) [\(level)] [\(system)] [\(file).\(function):\(line)] \(message)"
+    if !additionalContext.isEmpty {
+      description += "\ncontext: \(additionalContext.description)"
+    }
+    return description
   }
 }
 
@@ -68,7 +82,8 @@ extension SupabaseLogger {
     message: @autoclosure () -> String,
     fileID: StaticString = #fileID,
     function: StaticString = #function,
-    line: UInt = #line
+    line: UInt = #line,
+    additionalContext: JSONObject = [:]
   ) {
     let system = "\(fileID)".split(separator: "/").first ?? ""
 
@@ -80,7 +95,11 @@ extension SupabaseLogger {
         fileID: "\(fileID)",
         function: "\(function)",
         line: line,
-        timestamp: Date().timeIntervalSince1970
+        timestamp: Date().timeIntervalSince1970,
+        additionalContext: additionalContext.merging(
+          SupabaseLoggerTaskLocal.additionalContext,
+          uniquingKeysWith: { _, new in new }
+        )
       )
     )
   }
@@ -90,38 +109,89 @@ extension SupabaseLogger {
     _ message: @autoclosure () -> String,
     fileID: StaticString = #fileID,
     function: StaticString = #function,
-    line: UInt = #line
+    line: UInt = #line,
+    additionalContext: JSONObject = [:]
   ) {
-    log(.verbose, message: message(), fileID: fileID, function: function, line: line)
+    log(
+      .verbose,
+      message: message(),
+      fileID: fileID,
+      function: function,
+      line: line,
+      additionalContext: additionalContext
+    )
   }
 
   @inlinable
   public func debug(
     _ message: @autoclosure () -> String,
     fileID: StaticString = #fileID,
     function: StaticString = #function,
-    line: UInt = #line
+    line: UInt = #line,
+    additionalContext: JSONObject = [:]
   ) {
-    log(.debug, message: message(), fileID: fileID, function: function, line: line)
+    log(
+      .debug,
+      message: message(),
+      fileID: fileID,
+      function: function,
+      line: line,
+      additionalContext: additionalContext
+    )
   }
 
   @inlinable
   public func warning(
     _ message: @autoclosure () -> String,
     fileID: StaticString = #fileID,
     function: StaticString = #function,
-    line: UInt = #line
+    line: UInt = #line,
+    additionalContext: JSONObject = [:]
   ) {
-    log(.warning, message: message(), fileID: fileID, function: function, line: line)
+    log(
+      .warning,
+      message: message(),
+      fileID: fileID,
+      function: function,
+      line: line,
+      additionalContext: additionalContext
+    )
   }
 
   @inlinable
   public func error(
     _ message: @autoclosure () -> String,
     fileID: StaticString = #fileID,
     function: StaticString = #function,
-    line: UInt = #line
+    line: UInt = #line,
+    additionalContext: JSONObject = [:]
   ) {
-    log(.error, message: message(), fileID: fileID, function: function, line: line)
+    log(
+      .error,
+      message: message(),
+      fileID: fileID,
+      function: function,
+      line: line,
+      additionalContext: additionalContext
+    )
+  }
+}
+
+@inlinable
+package func trace<R>(
+  using logger: (any SupabaseLogger)?,
+  @_inheritActorContext _ operation: @Sendable () async throws -> R,
+  fileID: StaticString = #fileID,
+  function: StaticString = #function,
+  line: UInt = #line
+) async rethrows -> R {
+  logger?.debug("begin", fileID: fileID, function: function, line: line)
+  defer { logger?.debug("end", fileID: fileID, function: function, line: line) }
+
+  do {
+    return try await operation()
+  } catch {
+    logger?.debug("error: \(error)", fileID: fileID, function: function, line: line)
+    throw error
   }
 }
diff --git a/Sources/_Helpers/TaskLocalHelpers.swift b/Sources/_Helpers/TaskLocalHelpers.swift
