From 2ab6d60df35a942d06bb7f33af1d442b7b5c59dc Mon Sep 17 00:00:00 2001 From: Mike Keesey Date: Mon, 23 Mar 2026 18:38:05 +0000 Subject: [PATCH] Allow usage of stale-but-valid ingestion auth tokens The Kusto ingestion client periodically (every 1h) attempts to call a management command on the destination Kusto cluster to refresh a token used to authenticate to Kusto-cluster owned storage. On occasion, we experience issues where our Kusto cluster's control plane has brief unavailability for several minutes at a time - this can completely block ingestion from clients for long periods of time, despite the tokens themselves having much longer lifetimes and the destination storage remaining available. Instead of failing fast when new tokens cannot be minted, this PR adds a mechanism to reuse old-but-valid tokens when new tokens can't be created. This should allow these clients to continue uploading to storage while the Kusto cluster is temporarily unavailable. This also adds a throttling mechanism where we do not continually hammer the destination Kusto cluster during token expirations, creating more of a thundering herd when the Kusto control plane is already having issues. --- azkustoingest/internal/resources/resources.go | 82 +++- .../internal/resources/resources_test.go | 355 +++++++++++++++++- 2 files changed, 418 insertions(+), 19 deletions(-) diff --git a/azkustoingest/internal/resources/resources.go b/azkustoingest/internal/resources/resources.go index 1c93ba82..14de3834 100644 --- a/azkustoingest/internal/resources/resources.go +++ b/azkustoingest/internal/resources/resources.go @@ -26,6 +26,11 @@ const ( defaultMultiplier = 2 retryCount = 4 fetchInterval = 1 * time.Hour + authRefreshInterval = 1 * time.Hour + authNoRefreshTimeout = 10 * time.Hour + + authRetryIntervalWithCachedToken = 15 * time.Minute + authRetryIntervalWithoutCachedToken = 3 * time.Second ) // mgmter is a private interface that allows us to write hermetic tests against the azkustodata.Client.Mgmt() method. @@ -116,15 +121,19 @@ type ResourcesManager interface { // Manager manages Kusto resources. type Manager struct { - client mgmter - done chan struct{} - resources atomic.Value // Stores Ingestion - lastFetchTime atomic.Value // Stores time.Time - kustoToken token - authTokenCacheExpiration time.Time - authLock sync.Mutex - fetchLock sync.Mutex - rankedStorageAccount *RankedStorageAccountSet + client mgmter + done chan struct{} + resources atomic.Value // Stores Ingestion + lastFetchTime atomic.Value // Stores time.Time + kustoToken token + hasAuthToken bool + authTokenRefreshTime time.Time + authTokenValidUntil time.Time + nextAuthRefreshAttempt time.Time + authLock sync.Mutex + fetchLock sync.Mutex + rankedStorageAccount *RankedStorageAccountSet + now func() time.Time } var _ ResourcesManager = (*Manager)(nil) @@ -134,8 +143,9 @@ func New(client mgmter) (*Manager, error) { m := &Manager{client: client, done: make(chan struct{}), rankedStorageAccount: newDefaultRankedStorageAccountSet()} m.authLock = sync.Mutex{} m.fetchLock = sync.Mutex{} + m.now = func() time.Time { return time.Now().UTC() } - m.authTokenCacheExpiration = time.Now().UTC() + m.authTokenRefreshTime = m.now() go m.renewResources() return m, nil @@ -180,10 +190,44 @@ func (m *Manager) renewResources() { func (m *Manager) AuthContext(ctx context.Context) (string, error) { m.authLock.Lock() defer m.authLock.Unlock() - if m.authTokenCacheExpiration.After(time.Now().UTC()) { + now := m.currentTime() + + if m.hasAuthToken && m.authTokenRefreshTime.After(now) { return m.kustoToken.AuthContext, nil } + if m.nextAuthRefreshAttempt.After(now) { + if m.hasAuthToken && m.authTokenValidUntil.After(now) { + return m.kustoToken.AuthContext, nil + } + + if m.nextAuthRefreshAttempt.Sub(now) <= authRetryIntervalWithoutCachedToken { + return "", fmt.Errorf("problem getting authorization context from Kusto via Mgmt: refresh is temporarily throttled") + } + } + + tk, err := m.refreshAuthContext(ctx) + now = m.currentTime() + if err != nil { + if m.hasAuthToken && m.authTokenValidUntil.After(now) { + m.nextAuthRefreshAttempt = now.Add(authRetryIntervalWithCachedToken) + return m.kustoToken.AuthContext, nil + } + + m.nextAuthRefreshAttempt = now.Add(authRetryIntervalWithoutCachedToken) + return "", fmt.Errorf("problem getting authorization context from Kusto via Mgmt: %w", err) + } + + m.kustoToken = tk + m.hasAuthToken = true + m.authTokenRefreshTime = now.Add(authRefreshInterval) + m.authTokenValidUntil = now.Add(authNoRefreshTimeout) + m.nextAuthRefreshAttempt = time.Time{} + + return tk.AuthContext, nil +} + +func (m *Manager) refreshAuthContext(ctx context.Context) (token, error) { var dataset v1.Dataset retryCtx := backoff.WithContext(initBackoff(), ctx) err := backoff.Retry(func() error { @@ -202,24 +246,26 @@ func (m *Manager) AuthContext(ctx context.Context) (string, error) { }, retryCtx) if err != nil { - return "", fmt.Errorf("problem getting authorization context from Kusto via Mgmt: %s", err) + return token{}, err } tokens, err := query.ToStructs[token](dataset) if err != nil { - return "", err + return token{}, err } if tokens == nil { - return "", fmt.Errorf("call for AuthContext returned no Rows") + return token{}, fmt.Errorf("call for AuthContext returned no Rows") } if len(tokens) != 1 { - return "", fmt.Errorf("call for AuthContext returned more than 1 Row") + return token{}, fmt.Errorf("call for AuthContext returned more than 1 Row") } - m.kustoToken = tokens[0] - m.authTokenCacheExpiration = time.Now().UTC().Add(time.Hour) - return tokens[0].AuthContext, nil + return tokens[0], nil +} + +func (m *Manager) currentTime() time.Time { + return m.now().UTC() } // ingestResc represents a kusto Mgmt() record about a resource diff --git a/azkustoingest/internal/resources/resources_test.go b/azkustoingest/internal/resources/resources_test.go index 701ace9f..1722db25 100644 --- a/azkustoingest/internal/resources/resources_test.go +++ b/azkustoingest/internal/resources/resources_test.go @@ -2,8 +2,14 @@ package resources import ( "context" + goErrors "errors" + "fmt" + "sync" "testing" + "time" + "github.com/Azure/azure-kusto-go/azkustodata" + dataErrors "github.com/Azure/azure-kusto-go/azkustodata/errors" v1 "github.com/Azure/azure-kusto-go/azkustodata/query/v1" "github.com/stretchr/testify/assert" @@ -12,6 +18,79 @@ import ( "github.com/Azure/azure-kusto-go/azkustodata/value" ) +type authResponse struct { + token string + err error + multiRows bool +} + +type authSequenceMgmt struct { + mu sync.Mutex + responses []authResponse + calls int +} + +func (a *authSequenceMgmt) Mgmt(_ context.Context, _ string, statement azkustodata.Statement, _ ...azkustodata.QueryOption) (v1.Dataset, error) { + if statement.String() != ".get kusto identity token" { + return nil, fmt.Errorf("unexpected statement: %s", statement.String()) + } + + a.mu.Lock() + defer a.mu.Unlock() + + a.calls++ + if len(a.responses) == 0 { + return nil, goErrors.New("no responses configured") + } + + idx := a.calls - 1 + if idx >= len(a.responses) { + idx = len(a.responses) - 1 + } + + resp := a.responses[idx] + if resp.err != nil { + return nil, resp.err + } + + if resp.multiRows { + return multiRowAuthDataset(resp.token) + } + + return authDataset(resp.token) +} + +func (a *authSequenceMgmt) callCount() int { + a.mu.Lock() + defer a.mu.Unlock() + return a.calls +} + +func authDataset(tok string) (v1.Dataset, error) { + return v1.NewDataset(context.Background(), dataErrors.OpMgmt, v1.V1{Tables: []v1.RawTable{{ + TableName: "Table", + Columns: []v1.RawColumn{{ + ColumnName: "AuthorizationContext", + ColumnType: string(types.String), + }}, + Rows: []v1.RawRow{{Row: []interface{}{tok}}}, + }}}) +} + +func multiRowAuthDataset(tok string) (v1.Dataset, error) { + return v1.NewDataset(context.Background(), dataErrors.OpMgmt, v1.V1{Tables: []v1.RawTable{{ + TableName: "Table", + Columns: []v1.RawColumn{{ + ColumnName: "AuthorizationContext", + ColumnType: string(types.String), + }}, + Rows: []v1.RawRow{ + {Row: []interface{}{tok}}, + {Row: []interface{}{tok + "-extra"}}, + }, + }}}) +} + func TestParse(t *testing.T) { t.Parallel() @@ -138,7 +217,7 @@ func TestAuthContext(t *testing.T) { test := test t.Run(test.desc, func(t *testing.T) { t.Parallel() - manager := &Manager{client: test.fakeMgmt} + manager := &Manager{client: test.fakeMgmt, now: func() time.Time { return time.Now().UTC() }} got, err := manager.AuthContext(context.Background()) @@ -154,6 +233,280 @@ func TestAuthContext(t *testing.T) { } } +func TestAuthContextCachedTokenWithinRefreshInterval(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + seq := &authSequenceMgmt{responses: []authResponse{{token: "authtoken-1"}}} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + got, err := m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 1, seq.callCount()) + + // Within refresh interval - cached token returned, no Mgmt call. + now = now.Add(authRefreshInterval - time.Second) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 1, seq.callCount()) +} + +func TestAuthContextSuccessfulRefreshReplacesToken(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + seq := &authSequenceMgmt{responses: []authResponse{ + {token: "authtoken-1"}, + {token: "authtoken-2"}, + }} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + got, err := m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 1, seq.callCount()) + + // Exactly at authTokenRefreshTime - refresh triggered, token replaced. + now = now.Add(authRefreshInterval) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-2", got) + assert.Equal(t, 2, seq.callCount()) + + // Cached token-2 returned without Mgmt call. + now = now.Add(time.Minute) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-2", got) + assert.Equal(t, 2, seq.callCount()) +} + +func TestAuthContextRefreshFailureReturnsStaleAndThrottles(t *testing.T) { + t.Parallel() + + cases := []struct { + name string + failResp authResponse + }{ + {name: "mgmt_error", failResp: authResponse{err: goErrors.New("refresh failed")}}, + {name: "parse_error", failResp: authResponse{token: "ignored", multiRows: true}}, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + seq := &authSequenceMgmt{responses: []authResponse{ + {token: "authtoken-1"}, + tc.failResp, + }} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + got, err := m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + + // Past refresh interval - refresh attempted and fails; stale token returned. + now = now.Add(authRefreshInterval + time.Second) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 2, seq.callCount()) + + // Within long throttle - cached returned, no new Mgmt call. + now = now.Add(time.Minute) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 2, seq.callCount()) + }) + } +} + +func TestAuthContextExpiredStaleReturnsError(t *testing.T) { + t.Parallel() + + originalStart := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + now := originalStart + seq := &authSequenceMgmt{responses: []authResponse{ + {token: "authtoken-1"}, + {err: goErrors.New("refresh failed")}, + {err: goErrors.New("refresh failed again")}, + }} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + got, err := m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + + // Failed refresh well before validUntil - sets long throttle, returns stale. + now = now.Add(authRefreshInterval + time.Second) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 2, seq.callCount()) + + // Exactly at authTokenValidUntil - cached token no longer valid, refresh attempted and fails. + now = originalStart.Add(authNoRefreshTimeout) + got, err = m.AuthContext(context.Background()) + assert.Error(t, err) + assert.Empty(t, got) + assert.Equal(t, 3, seq.callCount()) +} + +func TestAuthContextFailureWithoutCachedTokenThrottlesAndRetries(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + seq := &authSequenceMgmt{responses: []authResponse{ + {err: goErrors.New("boom")}, + {token: "authtoken-1"}, + }} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + got, err := m.AuthContext(context.Background()) + assert.Error(t, err) + assert.Empty(t, got) + assert.Equal(t, 1, seq.callCount()) + + // Within short throttle window - throttled, no Mgmt call. + now = now.Add(authRetryIntervalWithoutCachedToken - time.Second) + got, err = m.AuthContext(context.Background()) + assert.Error(t, err) + assert.Empty(t, got) + assert.Equal(t, 1, seq.callCount()) + + // Exactly at nextAuthRefreshAttempt - retry permitted and succeeds. + now = now.Add(time.Second) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 2, seq.callCount()) +} + +// TestAuthContextExpiredTokenBypassesLongThrottle exercises the unique branch +// where a cached token expires while a long throttle is still active, forcing +// a refresh attempt despite the throttle. +func TestAuthContextExpiredTokenBypassesLongThrottle(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + seq := &authSequenceMgmt{responses: []authResponse{ + {token: "authtoken-1"}, + {err: goErrors.New("refresh failed")}, + {err: goErrors.New("refresh failed again")}, + {err: goErrors.New("refresh failed yet again")}, + {token: "authtoken-2"}, + }} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + got, err := m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + + // Refresh fails close to token expiry. Long throttle (15m) extends past + // authTokenValidUntil (10h) so it remains active after the token expires. + // authTokenValidUntil = 10:00:00, advance to 09:55:00 -> nextAuthRefreshAttempt = 10:10:00. + now = now.Add(authNoRefreshTimeout - 5*time.Minute) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-1", got) + assert.Equal(t, 2, seq.callCount()) + + // Token expired (>10:00:00) but long throttle still active (until 10:10:00). + // Refresh attempt is made anyway because cached token is no longer valid. + now = now.Add(7 * time.Minute) // 10:02:00 + got, err = m.AuthContext(context.Background()) + assert.Error(t, err) + assert.Empty(t, got) + assert.Equal(t, 3, seq.callCount()) + + // After failed bypass, nextAuthRefreshAttempt is now short (3s). Immediate retry is throttled. + now = now.Add(time.Second) + got, err = m.AuthContext(context.Background()) + assert.Error(t, err) + assert.Empty(t, got) + assert.Equal(t, 3, seq.callCount()) + + // Past short throttle, refresh attempted again and fails. + now = now.Add(authRetryIntervalWithoutCachedToken) + got, err = m.AuthContext(context.Background()) + assert.Error(t, err) + assert.Empty(t, got) + assert.Equal(t, 4, seq.callCount()) + + // After short throttle elapses, refresh succeeds. + now = now.Add(authRetryIntervalWithoutCachedToken) + got, err = m.AuthContext(context.Background()) + assert.NoError(t, err) + assert.Equal(t, "authtoken-2", got) + assert.Equal(t, 5, seq.callCount()) +} + +func TestAuthContextConcurrentCallersTriggerSingleRefresh(t *testing.T) { + t.Parallel() + + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + seq := &authSequenceMgmt{responses: []authResponse{{token: "authtoken-1"}}} + m := &Manager{ + client: seq, + now: func() time.Time { return now }, + } + + const workers = 32 + start := make(chan struct{}) + var wg sync.WaitGroup + + errs := make(chan error, workers) + tokens := make(chan string, workers) + + for i := 0; i < workers; i++ { + wg.Add(1) + go func() { + defer wg.Done() + <-start + tk, err := m.AuthContext(context.Background()) + errs <- err + tokens <- tk + }() + } + + close(start) + wg.Wait() + close(errs) + close(tokens) + + for err := range errs { + assert.NoError(t, err) + } + for tk := range tokens { + assert.Equal(t, "authtoken-1", tk) + } + + assert.Equal(t, 1, seq.callCount()) +} + func mustParse(s string) *URI { u, err := Parse(s) if err != nil {