Only request once if cache is still valid

lib/livebook/teams/requests.ex

@@ -225,17 +225,21 @@ defmodule Livebook.Teams.Requests do
   @doc """
   Send a request to Livebook Team API to get the user information from given access token.
   """
-  @spec get_user_info(Team.t(), String.t()) :: api_result() | :econnrefused
-  def get_user_info(team, access_token) do
+  @spec get_user_info(Team.t(), String.t(), boolean()) :: api_result() | :econnrefused
+  def get_user_info(team, access_token, valid_cache?) do
     req = build_req(team)
     params = %{access_token: access_token}
 
-    case Req.get(req, url: "/api/v1/org/identity", params: params) do
-      {:error, %Req.TransportError{reason: :econnrefused}} ->
-        :econnrefused
+    fun = fn
+      _, %Req.TransportError{reason: :econnrefused} -> not valid_cache?
+      _, %Req.Response{status: status} when status in [408, 429, 500, 502, 503, 504] -> true
+      _, %Req.TransportError{reason: reason} when reason in [:timeout, :closed] -> true
+      _, _ -> false
+    end
 
-      otherwise ->
-        handle_response(otherwise)
+    case Req.get(req, url: "/api/v1/org/identity", params: params, retry: fun) do
+      {:error, %Req.TransportError{reason: :econnrefused}} -> :econnrefused
+      otherwise -> handle_response(otherwise)
     end
   end
 
@@ -330,17 +334,6 @@ defmodule Livebook.Teams.Requests do
     |> Req.Request.put_new_header("x-lb-version", Livebook.Config.app_version())
     |> Livebook.Utils.req_attach_defaults()
     |> add_team_auth(team)
-    |> put_test_req_options()
-  end
-
-  if Mix.env() == :test do
-    defp put_test_req_options(req) do
-      Req.Request.merge_options(req, retry: false)
-    end
-  else
-    defp put_test_req_options(req) do
-      req
-    end
   end
 
   defp add_team_auth(req, nil), do: req

lib/livebook/zta/livebook_teams.ex

@@ -69,7 +69,7 @@ defmodule Livebook.ZTA.LivebookTeams do
 
   defp handle_request(name, conn, team, %{"teams_identity" => _, "code" => code}) do
     with {:ok, access_token} <- retrieve_access_token(team, code),
-         {:ok, payload} <- Teams.Requests.get_user_info(team, access_token) do
+         {:ok, payload} <- Teams.Requests.get_user_info(team, access_token, false) do
       metadata = build_metadata(team.id, payload)
       exp = System.os_time(:second) + 3 * 3600
       :ets.insert(name, {access_token, {exp, metadata}})
@@ -183,30 +183,33 @@ defmodule Livebook.ZTA.LivebookTeams do
   defp validate_access_token(name, conn, team, access_token) do
     node = get_session(conn, :livebook_teams_metadata_node)
 
-    case Teams.Requests.get_user_info(team, access_token) do
+    entry =
+      try do
+        :erpc.call(node, :ets, :lookup_element, [name, access_token, 2, nil])
+      catch
+        _, _ -> nil
+      end
+
+    valid_cache? = valid_cache?(entry)
+
+    case Teams.Requests.get_user_info(team, access_token, valid_cache?) do
       {:ok, payload} ->
         {conn, build_metadata(team.id, payload)}
 
       :econnrefused ->
-        entry =
-          try do
-            :erpc.call(node, :ets, :lookup_element, [name, access_token, 2, nil])
-          catch
-            _, _ -> nil
-          end
-
-        case {System.os_time(:second), entry} do
-          {current_timestamp, {exp, metadata}} when current_timestamp <= exp ->
-            {conn, metadata}
-
-          {_, entry} ->
-            entry && :erpc.call(node, :ets, :delete, [name, access_token])
-
-            {conn
-             |> put_status(:service_unavailable)
-             |> put_view(LivebookWeb.ErrorHTML)
-             |> render("503.html")
-             |> halt(), nil}
+        # We double-checked because the response may contain latency, 
+        # so the timestamp must be revalidated to ensure it hasn't expired.
+        if valid_cache?(entry) do
+          {_, metadata} = entry
+          {conn, metadata}
+        else
+          entry && :erpc.call(node, :ets, :delete, [name, access_token])
+
+          {conn
+           |> put_status(:service_unavailable)
+           |> put_view(LivebookWeb.ErrorHTML)
+           |> render("503.html")
+           |> halt(), nil}
         end
 
       _otherwise ->
@@ -220,6 +223,9 @@ defmodule Livebook.ZTA.LivebookTeams do
     end
   end
 
+  defp valid_cache?({exp, _}), do: System.os_time(:second) <= exp
+  defp valid_cache?(_entry), do: false
+
   @doc """
   Returns the user metadata from given payload.
   """

test/livebook_teams/zta/livebook_teams_test.exs

@@ -134,6 +134,7 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
 
     alias Livebook.ZTA.LivebookTeams
 
+    @moduletag :capture_log
     @moduletag teams_for: :agent
     setup :teams
 
@@ -142,40 +143,37 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
 
     setup :livebook_teams_auth
 
-    test "uses cached version of the identity payload",
-         %{conn: conn, test: test, node: node, code: code} do
+    test "uses cached version of the identity payload", %{conn: conn, test: test} = ctx do
+      Application.put_env(:livebook, :teams_url, "http://localhost:1234")
+
       id = conn.assigns.current_user.id
       access_token = get_session(conn, :livebook_teams_access_token)
       groups = [%{"provider_id" => "1", "group_name" => "Foo"}]
-
-      # simulate the Teams API is down
-      url = Livebook.Config.teams_url()
-      Application.put_env(:livebook, :teams_url, "http://localhost:1234")
+      node = get_session(conn, :livebook_teams_metadata_node)
 
       # update the groups, but doesn't return because Livebook is using the cached one
-      TeamsRPC.update_user_info_groups(node, code, groups)
-      assert {_, %{id: ^id, groups: []}} = LivebookTeams.authenticate(test, conn, [])
+      TeamsRPC.update_user_info_groups(ctx.node, ctx.code, groups)
+
+      # shouldn't retry the request
+      assert ExUnit.CaptureLog.capture_log(fn ->
+               assert {_, %{id: ^id, groups: []}} = LivebookTeams.authenticate(test, conn, [])
+             end) == ""
 
       # simulate if the token already expired
       exp = System.os_time(:second) - 5 * 60
-      metadata_node = get_session(conn, :livebook_teams_metadata_node)
-
-      {_, metadata} =
-        :erpc.call(metadata_node, :ets, :lookup_element, [test, access_token, 2, nil])
-
-      :erpc.call(metadata_node, :ets, :insert, [test, {access_token, {exp, metadata}}])
-
-      # now it should return status 503
-      assert {%{status: 503, halted: true, resp_body: body}, nil} =
-               LivebookTeams.authenticate(test, conn, [])
+      {_, metadata} = :erpc.call(node, :ets, :lookup_element, [test, access_token, 2, nil])
+      :erpc.call(node, :ets, :insert, [test, {access_token, {exp, metadata}}])
 
-      assert body =~ "The server is currently down or under maintenance"
+      # now it should retry to request to Teams and return status 503
+      assert ExUnit.CaptureLog.capture_log(fn ->
+               assert {%{status: 503, halted: true, resp_body: body}, nil} =
+                        LivebookTeams.authenticate(test, conn, [])
 
-      # still show 503 error page because Teams isn't up yet
-      assert {%{status: 503, halted: true}, nil} = LivebookTeams.authenticate(test, conn, [])
+               assert body =~ "The server is currently down or under maintenance"
+             end) =~ "retry: got exception, will retry in"
 
       # now gets the updated userinfo from Teams
-      Application.put_env(:livebook, :teams_url, url)
+      Application.put_env(:livebook, :teams_url, TeamsServer.url())
       assert {_conn, %{id: ^id, groups: ^groups}} = LivebookTeams.authenticate(test, conn, [])
     end
   end