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 {