From 6e207dab536dbc983b46455a36815ad44b50e3aa Mon Sep 17 00:00:00 2001 From: Prentice Roosevelt Date: Wed, 7 Oct 2026 14:22:37 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20Refresh=20Kinesis=20client=20cre?= =?UTF-8?q?dentials=20when=20using=20task=20role?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Kinesis pipeline built its AWS client once in init/2 and reused it forever. With task role credentials (ECS task role, EKS IRSA, instance profile) the session token expires (1h for IRSA), so every PutRecords failed with "The security token included in the request is expired" until the pipeline restarted, even though aws_credentials had already refreshed the credentials. Rebuild the client from aws_credentials at most once a minute when the sink uses task role. aws_credentials refreshes credentials 5 minutes before expiry, so the client always carries a valid token. Sinks with explicit access keys are unchanged. Refs sequinstream/sequin#2078 Co-Authored-By: Claude Opus 5.5 --- lib/sequin/runtime/kinesis_pipeline.ex | 50 ++++++- .../kinesis_pipeline_task_role_test.exs | 138 ++++++++++++++++++ 2 files changed, 185 insertions(+), 3 deletions(-) create mode 100644 test/sequin/kinesis_pipeline_task_role_test.exs diff --git a/lib/sequin/runtime/kinesis_pipeline.ex b/lib/sequin/runtime/kinesis_pipeline.ex index 3c435adc1..69c68c24a 100644 --- a/lib/sequin/runtime/kinesis_pipeline.ex +++ b/lib/sequin/runtime/kinesis_pipeline.ex @@ -9,9 +9,15 @@ defmodule Sequin.Runtime.KinesisPipeline do alias Sequin.Runtime.Routing alias Sequin.Runtime.SinkPipeline + require Logger + # Kinesis PartitionKey has a maximum of 256 bytes @max_partition_key_bytes 256 + # aws_credentials refreshes temporary credentials 5 minutes before they expire, + # so rebuilding the client at most once a minute always picks up a valid token. + @default_task_role_client_refresh_seconds 60 + @impl SinkPipeline def init(context, _opts) do %{consumer: consumer, test_pid: test_pid} = context @@ -20,7 +26,9 @@ defmodule Sequin.Runtime.KinesisPipeline do case KinesisSink.aws_client(consumer.sink) do {:ok, client} -> - Map.put(context, :kinesis_client, client) + context + |> Map.put(:kinesis_client, client) + |> Map.put(:kinesis_client_built_at, now_seconds()) {:error, reason} -> raise "Failed to initialize Kinesis client: #{inspect(reason)}" @@ -54,16 +62,20 @@ defmodule Sequin.Runtime.KinesisPipeline do def handle_batch(:default, messages, %{batch_key: stream_arn}, context) do %{ consumer: %SinkConsumer{} = consumer, - kinesis_client: kinesis_client, test_pid: test_pid } = context setup_allowances(test_pid) + # With task role (ECS task role, IRSA, instance profile) the credentials are + # temporary. aws_credentials refreshes them before expiry, but the client built + # in init/2 keeps the original token, so rebuild it periodically. + context = maybe_refresh_client(context) + records = Enum.map(messages, fn message -> build_kinesis_record(consumer, message) end) - case Kinesis.put_records(kinesis_client, stream_arn, records) do + case Kinesis.put_records(context.kinesis_client, stream_arn, records) do :ok -> {:ok, messages, context} {:error, error} -> {:error, error} end @@ -82,6 +94,38 @@ defmodule Sequin.Runtime.KinesisPipeline do } end + defp maybe_refresh_client(%{consumer: %SinkConsumer{sink: %KinesisSink{use_task_role: true} = sink}} = context) do + built_at = Map.get(context, :kinesis_client_built_at) + + # monotonic time can be negative, so a missing timestamp is checked with is_nil + if is_nil(built_at) or now_seconds() - built_at >= task_role_client_refresh_seconds() do + case KinesisSink.aws_client(sink) do + {:ok, client} -> + context + |> Map.put(:kinesis_client, client) + |> Map.put(:kinesis_client_built_at, now_seconds()) + + {:error, reason} -> + Logger.warning("[KinesisPipeline] Failed to refresh AWS client for task role: #{inspect(reason)}") + context + end + else + context + end + end + + defp maybe_refresh_client(context), do: context + + defp now_seconds, do: System.monotonic_time(:second) + + defp task_role_client_refresh_seconds do + Keyword.get( + Application.get_env(:sequin, __MODULE__, []), + :task_role_client_refresh_seconds, + @default_task_role_client_refresh_seconds + ) + end + defp setup_allowances(nil), do: :ok defp setup_allowances(test_pid) do diff --git a/test/sequin/kinesis_pipeline_task_role_test.exs b/test/sequin/kinesis_pipeline_task_role_test.exs new file mode 100644 index 000000000..64cac097b --- /dev/null +++ b/test/sequin/kinesis_pipeline_task_role_test.exs @@ -0,0 +1,138 @@ +defmodule Sequin.Runtime.KinesisPipelineTaskRoleTest do + # Not async: the refresh interval is set through the application env + use Sequin.DataCase, async: false + + alias Sequin.Aws.HttpClient + alias Sequin.AwsMock + alias Sequin.Consumers + alias Sequin.Consumers.ConsumerEvent + alias Sequin.Databases.ConnectionCache + alias Sequin.Factory.AccountsFactory + alias Sequin.Factory.AwsFactory + alias Sequin.Factory.CharacterFactory + alias Sequin.Factory.ConsumersFactory + alias Sequin.Factory.DatabasesFactory + alias Sequin.Factory.ReplicationFactory + alias Sequin.Runtime.KinesisPipeline + alias Sequin.Runtime.SinkPipeline + + setup do + account = AccountsFactory.insert_account!() + + postgres_database = + DatabasesFactory.insert_configured_postgres_database!(account_id: account.id, tables: :character_tables) + + ConnectionCache.cache_connection(postgres_database, Sequin.Repo) + + replication = + ReplicationFactory.insert_postgres_replication!( + account_id: account.id, + postgres_database_id: postgres_database.id + ) + + consumer = + ConsumersFactory.insert_sink_consumer!( + account_id: account.id, + type: :kinesis, + batch_size: 10, + replication_slot_id: replication.id + ) + + {:ok, consumer} = + Consumers.update_sink_consumer(consumer, %{ + sink: %{ + type: :kinesis, + stream_arn: "arn:aws:kinesis:us-east-1:123456789012:stream/test-stream", + region: "us-east-1", + use_task_role: true, + routing_mode: :static + } + }) + + # Each call to get_client returns a new set of temporary credentials, + # as aws_credentials does after refreshing an expiring token + {:ok, counter} = Agent.start_link(fn -> 0 end) + + Mox.stub(AwsMock, :get_client, fn "us-east-1" -> + n = Agent.get_and_update(counter, &{&1 + 1, &1 + 1}) + + {:ok, + "task-role-key-#{n}" + |> AWS.Client.create("task-role-secret-#{n}", "us-east-1") + |> Map.put(:session_token, "task-role-token-#{n}")} + end) + + test_pid = self() + + Req.Test.stub(HttpClient, fn conn -> + {_key, auth} = Enum.find(conn.req_headers, fn {key, _value} -> key == "authorization" end) + [_, access_key] = Regex.run(~r/Credential=([^\/]+)\//, auth) + send(test_pid, {:kinesis_access_key, access_key}) + Req.Test.json(conn, AwsFactory.kinesis_put_records_response_success()) + end) + + previous_env = Application.get_env(:sequin, KinesisPipeline) + on_exit(fn -> restore_env(previous_env) end) + + {:ok, %{consumer: consumer}} + end + + test "rebuilds the client with fresh task role credentials before delivering", %{consumer: consumer} do + Application.put_env(:sequin, KinesisPipeline, task_role_client_refresh_seconds: 0) + + start_pipeline!(consumer) + + ref1 = send_test_event(consumer) + assert_receive {:ack, ^ref1, [%{data: %ConsumerEvent{}}], []}, 1_000 + assert_receive {:kinesis_access_key, key1}, 1_000 + + ref2 = send_test_event(consumer) + assert_receive {:ack, ^ref2, [%{data: %ConsumerEvent{}}], []}, 1_000 + assert_receive {:kinesis_access_key, key2}, 1_000 + + # The credentials obtained on init are replaced on every batch + assert key1 != key2 + end + + test "delivers with task role credentials while the client is still fresh", %{consumer: consumer} do + Application.put_env(:sequin, KinesisPipeline, task_role_client_refresh_seconds: 3600) + + start_pipeline!(consumer) + + ref = send_test_batch(consumer, [deliverable_event(consumer), deliverable_event(consumer)]) + assert_receive {:ack, ^ref, [_, _], []}, 1_000 + assert_receive {:kinesis_access_key, "task-role-key-" <> _}, 1_000 + end + + defp restore_env(nil), do: Application.delete_env(:sequin, KinesisPipeline) + defp restore_env(env), do: Application.put_env(:sequin, KinesisPipeline, env) + + defp deliverable_event(consumer) do + [consumer_id: consumer.id] + |> ConsumersFactory.insert_deliverable_consumer_event!() + |> Map.put(:data, ConsumersFactory.consumer_event_data(record: CharacterFactory.character_attrs())) + end + + defp start_pipeline!(consumer) do + start_supervised!( + {SinkPipeline, + [ + consumer_id: consumer.id, + producer: Broadway.DummyProducer, + test_pid: self() + ]} + ) + end + + defp send_test_event(consumer) do + Broadway.test_message(broadway(consumer), deliverable_event(consumer), metadata: %{topic: "test", headers: []}) + end + + defp send_test_batch(consumer, events) do + Broadway.test_batch(broadway(consumer), events, metadata: %{topic: "test", headers: []}) + end + + defp broadway(consumer) do + SinkPipeline.via_tuple(consumer.id) + end +end