From 025bfc5aa0f1d21bbd1e2f18e6a40b31646aa03a Mon Sep 17 00:00:00 2001 From: Ayana Sarkar Date: Thu, 17 Sep 2026 02:03:11 +0530 Subject: [PATCH 1/4] fix(ratelimiter): set TTL on initial olric Put to prevent immortal keys Store.Get created new/reset rate-limit keys with a Put that set no TTL, then applied the TTL in a separate follow-up write. Any of the following left that window open and the key permanently immortal, since nothing else in the store ever re-applies a TTL to a key that already lacks one: - rate.Period == 0 for a limiter tier, so the TTL-setting write never ran - the follow-up write's ErrKeyNotFound error was silently ignored - a crash/restart landed between the two writes Fold the TTL into the same Put that creates/resets the key instead, removing the second write (and its round trip) entirely. Fixes #1571 Signed-off-by: Ayana Sarkar --- .../ratelimiter/olric_store.go | 37 ++++++++++++------- 1 file changed, 24 insertions(+), 13 deletions(-) diff --git a/src/invocation-plane-services/ratelimiter/olric_store.go b/src/invocation-plane-services/ratelimiter/olric_store.go index 446ec215fd..17a82c54fb 100644 --- a/src/invocation-plane-services/ratelimiter/olric_store.go +++ b/src/invocation-plane-services/ratelimiter/olric_store.go @@ -99,7 +99,13 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim value := 1 // initialize the value to 1 if key doesn't exist if errors.Is(err, olric.ErrKeyNotFound) { - err = store.dmap.Put(ctx, fullKey, 1) + // Set the TTL on this same write instead of as a follow-up Put (see #1571): + // a key that is ever written without a TTL is never given one later, so + // creating it with no expiry and only setting the TTL in a second write + // left a window - rate.Period == 0, a silently-ignored ErrKeyNotFound on + // the follow-up write, or a crash/restart between the two writes - where + // the key would live forever and leak memory. + err = store.dmap.Put(ctx, fullKey, value, newKeyPutOptions(rate)...) if err != nil { return limiter.Context{}, err } @@ -115,8 +121,9 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim return limiter.Context{}, err } } else { - // when ttl is 0, reset the value to 1 - err = store.dmap.Put(ctx, fullKey, 1) + // when ttl is 0, reset the value to 1 and (re)apply the TTL on this + // same write so a key whose TTL already expired can't go immortal. + err = store.dmap.Put(ctx, fullKey, value, newKeyPutOptions(rate)...) if err != nil { zap.L().Error("Failed to update value", zap.Error(err)) return limiter.Context{}, err @@ -128,16 +135,6 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim now := time.Now() expiration := now.Add(rate.Period) if value == 1 { - if rate.Period.Milliseconds() > 0 { - ttlOption := olric.PX(rate.Period) - err = store.dmap.Put(ctx, fullKey, value, ttlOption) - - // ignore key not found error - if err != nil && !errors.Is(err, olric.ErrKeyNotFound) { - zap.L().Error("Failed to set expiration", zap.Error(err)) - return limiter.Context{}, err - } - } return common.GetContextFromState(now, rate, expiration, int64(value)), nil } @@ -153,6 +150,20 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim return common.GetContextFromState(now, rate, expiration, int64(value)), nil } +// newKeyPutOptions returns the Olric put options used whenever a key is +// (re)created with value 1. Folding the TTL into that same Put call - rather +// than applying it in a separate follow-up write - is what closes #1571: +// nothing in this store ever re-applies a TTL to a key that already lacks +// one, so the key must get its TTL atomically with the write that creates it. +// When rate.Period is zero (an unbounded limiter tier), no TTL is applied, +// matching the previous behavior for that case. +func newKeyPutOptions(rate limiter.Rate) []olric.PutOption { + if rate.Period.Milliseconds() > 0 { + return []olric.PutOption{olric.PX(rate.Period)} + } + return nil +} + // Peek returns the limit for the given identifier, without modification on current values. // NOT USED func (store *Store) Peek(ctx context.Context, key string, rate limiter.Rate) (limiter.Context, error) { From a4c6f59b3cfbdc78a67d144ac057e3cfe5adf06c Mon Sep 17 00:00:00 2001 From: Ayana Sarkar Date: Thu, 17 Sep 2026 02:33:41 +0530 Subject: [PATCH 2/4] test(ratelimiter): cover TTL-on-creating-write fix for #1571 Signed-off-by: Ayana Sarkar --- .../ratelimiter/olric_store_test.go | 154 ++++++++++++++++++ 1 file changed, 154 insertions(+) create mode 100644 src/invocation-plane-services/ratelimiter/olric_store_test.go diff --git a/src/invocation-plane-services/ratelimiter/olric_store_test.go b/src/invocation-plane-services/ratelimiter/olric_store_test.go new file mode 100644 index 0000000000..56fdbdb89d --- /dev/null +++ b/src/invocation-plane-services/ratelimiter/olric_store_test.go @@ -0,0 +1,154 @@ +/* +SPDX-FileCopyrightText: Copyright (c) NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package ratelimiter + +import ( + "context" + "testing" + "time" + + "github.com/olric-data/olric" + "github.com/ulule/limiter/v3" +) + +// fakePutCall records a single call made to fakeDMap.Put. +type fakePutCall struct { + key string + value interface{} + numOptions int +} + +// fakeDMap is a minimal olric.DMap test double. Store.Get and Store.Reset +// only ever call Put, Get and Delete, so every other method panics if +// called - an unexpected dependency on them should fail the test loudly +// rather than silently return a zero value. +type fakeDMap struct { + getResult *olric.GetResponse + getErr error + + putCalls []fakePutCall + putErr error +} + +var _ olric.DMap = (*fakeDMap)(nil) + +func (f *fakeDMap) Name() string { return "test" } + +func (f *fakeDMap) Put(_ context.Context, key string, value interface{}, options ...olric.PutOption) error { + f.putCalls = append(f.putCalls, fakePutCall{key: key, value: value, numOptions: len(options)}) + return f.putErr +} + +func (f *fakeDMap) Get(context.Context, string) (*olric.GetResponse, error) { + return f.getResult, f.getErr +} + +func (f *fakeDMap) Delete(context.Context, ...string) (int, error) { panic("not implemented") } +func (f *fakeDMap) Incr(context.Context, string, int) (int, error) { panic("not implemented") } +func (f *fakeDMap) Decr(context.Context, string, int) (int, error) { panic("not implemented") } +func (f *fakeDMap) GetPut(context.Context, string, interface{}) (*olric.GetResponse, error) { + panic("not implemented") +} +func (f *fakeDMap) IncrByFloat(context.Context, string, float64) (float64, error) { + panic("not implemented") +} +func (f *fakeDMap) Expire(context.Context, string, time.Duration) error { panic("not implemented") } +func (f *fakeDMap) Lock(context.Context, string, time.Duration) (olric.LockContext, error) { + panic("not implemented") +} +func (f *fakeDMap) LockWithTimeout(context.Context, string, time.Duration, time.Duration) (olric.LockContext, error) { + panic("not implemented") +} +func (f *fakeDMap) Scan(context.Context, ...olric.ScanOption) (olric.Iterator, error) { + panic("not implemented") +} +func (f *fakeDMap) Destroy(context.Context) error { panic("not implemented") } +func (f *fakeDMap) Pipeline(...olric.PipelineOption) (*olric.DMapPipeline, error) { + panic("not implemented") +} +func (f *fakeDMap) Close(context.Context) error { panic("not implemented") } + +// TestGet_NewKey_SetsTTLOnCreatingWrite is the core regression test for +// #1571: a brand-new key (dmap.Get returns ErrKeyNotFound) must be created +// with exactly one Put call that already carries its TTL. Before the fix, +// this path wrote the key once with no TTL and only set the TTL in a +// second, separate Put - which is exactly the gap that produced immortal +// keys. +func TestGet_NewKey_SetsTTLOnCreatingWrite(t *testing.T) { + f := &fakeDMap{getErr: olric.ErrKeyNotFound} + store := &Store{Prefix: "test", dmap: f} + + if _, err := store.Get(context.Background(), "user-1", limiter.Rate{Period: time.Second, Limit: 10}); err != nil { + t.Fatalf("Get returned error: %v", err) + } + + if len(f.putCalls) != 1 { + t.Fatalf("expected exactly 1 Put call for a new key, got %d - a second follow-up write reopens the immortal-key gap from #1571", len(f.putCalls)) + } + if got := f.putCalls[0]; got.numOptions != 1 { + t.Fatalf("expected the creating Put to carry 1 TTL option, got %d - a new key must never be written without a TTL (#1571)", got.numOptions) + } + if f.putCalls[0].value != 1 { + t.Fatalf("expected new key value 1, got %v", f.putCalls[0].value) + } +} + +// TestGet_NewKey_UnboundedTier_NoTTLOption preserves the pre-existing +// behavior for an unbounded limiter tier (rate.Period == 0): still exactly +// one Put, but with no TTL option, matching what the code did before #1571 +// for this specific case. +func TestGet_NewKey_UnboundedTier_NoTTLOption(t *testing.T) { + f := &fakeDMap{getErr: olric.ErrKeyNotFound} + store := &Store{Prefix: "test", dmap: f} + + if _, err := store.Get(context.Background(), "user-1", limiter.Rate{Period: 0, Limit: 10}); err != nil { + t.Fatalf("Get returned error: %v", err) + } + + if len(f.putCalls) != 1 { + t.Fatalf("expected exactly 1 Put call, got %d", len(f.putCalls)) + } + if got := f.putCalls[0].numOptions; got != 0 { + t.Fatalf("expected no TTL option for an unbounded (Period == 0) tier, got %d options", got) + } +} + +// TestGet_ExpiredKey_ResetGetsTTLOnSameWrite covers the second branch #1571 +// identified: an existing key whose TTL has already lapsed (dmap.Get +// returns a nil result with no error) is reset to value 1. That reset must +// also carry the TTL on the same write, or the key goes immortal exactly +// like the new-key case above. +func TestGet_ExpiredKey_ResetGetsTTLOnSameWrite(t *testing.T) { + f := &fakeDMap{} // getResult == nil, getErr == nil: key exists but expired/untimed + store := &Store{Prefix: "test", dmap: f} + + if _, err := store.Get(context.Background(), "user-1", limiter.Rate{Period: time.Second, Limit: 10}); err != nil { + t.Fatalf("Get returned error: %v", err) + } + + if len(f.putCalls) != 1 { + t.Fatalf("expected exactly 1 Put call when resetting an expired key, got %d", len(f.putCalls)) + } + if got := f.putCalls[0]; got.numOptions != 1 { + t.Fatalf("expected the reset Put to carry 1 TTL option, got %d - a key whose TTL already lapsed must never be rewritten without one (#1571)", got.numOptions) + } +} + +func (f *fakeDMap) CompareAndSwap(context.Context, string, []byte, interface{}, ...olric.PutOption) (bool, *olric.GetResponse, error) { +panic("not implemented") +} From 6b118088acc83778f96649de4cb7538c77a0a1b7 Mon Sep 17 00:00:00 2001 From: Ayana Sarkar Date: Thu, 17 Sep 2026 03:12:50 +0530 Subject: [PATCH 3/4] fix(ratelimiter): round sub-millisecond periods up to Olric's PX precision Signed-off-by: Ayana Sarkar --- .../ratelimiter/olric_store.go | 8 ++++-- .../ratelimiter/olric_store_test.go | 26 ++++++++++++++++++- 2 files changed, 31 insertions(+), 3 deletions(-) diff --git a/src/invocation-plane-services/ratelimiter/olric_store.go b/src/invocation-plane-services/ratelimiter/olric_store.go index 17a82c54fb..dcdb138752 100644 --- a/src/invocation-plane-services/ratelimiter/olric_store.go +++ b/src/invocation-plane-services/ratelimiter/olric_store.go @@ -158,8 +158,12 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim // When rate.Period is zero (an unbounded limiter tier), no TTL is applied, // matching the previous behavior for that case. func newKeyPutOptions(rate limiter.Rate) []olric.PutOption { - if rate.Period.Milliseconds() > 0 { - return []olric.PutOption{olric.PX(rate.Period)} + if rate.Period > 0 { + ttl := rate.Period + if ttl < time.Millisecond { + ttl = time.Millisecond + } + return []olric.PutOption{olric.PX(ttl)} } return nil } diff --git a/src/invocation-plane-services/ratelimiter/olric_store_test.go b/src/invocation-plane-services/ratelimiter/olric_store_test.go index 56fdbdb89d..e7b8da120b 100644 --- a/src/invocation-plane-services/ratelimiter/olric_store_test.go +++ b/src/invocation-plane-services/ratelimiter/olric_store_test.go @@ -150,5 +150,29 @@ func TestGet_ExpiredKey_ResetGetsTTLOnSameWrite(t *testing.T) { } func (f *fakeDMap) CompareAndSwap(context.Context, string, []byte, interface{}, ...olric.PutOption) (bool, *olric.GetResponse, error) { -panic("not implemented") + panic("not implemented") +} + +// TestGet_NewKey_SubMillisecondPeriod_TTLRoundedUp covers a direct Store.Get +// caller supplying a period below Olric's one-millisecond PX precision. +// Before this fix, newKeyPutOptions treated any period under 1ms as +// unbounded and wrote no TTL at all - reopening the immortal-key gap from +// #1571 for non-configured (programmatic) callers. Configured rates +// (parseRates) only ever produce S/M/H/D-granularity periods, so this case +// cannot occur from config today, but Store.Get itself accepts any +// limiter.Rate. +func TestGet_NewKey_SubMillisecondPeriod_TTLRoundedUp(t *testing.T) { + f := &fakeDMap{getErr: olric.ErrKeyNotFound} + store := &Store{Prefix: "test", dmap: f} + + if _, err := store.Get(context.Background(), "user-1", limiter.Rate{Period: 500 * time.Microsecond, Limit: 10}); err != nil { + t.Fatalf("Get returned error: %v", err) + } + + if len(f.putCalls) != 1 { + t.Fatalf("expected exactly 1 Put call, got %d", len(f.putCalls)) + } + if got := f.putCalls[0].numOptions; got != 1 { + t.Fatalf("expected a TTL option even for a sub-millisecond period, got %d options - a positive period must never be silently treated as unbounded", got) + } } From 1b60689ce9d332f6d70aa76c0205897abb6170e4 Mon Sep 17 00:00:00 2001 From: Ayana Sarkar Date: Thu, 17 Sep 2026 03:41:07 +0530 Subject: [PATCH 4/4] fix(ratelimiter): treat result.TTL() as absolute epoch-ms, not a duration Signed-off-by: Ayana Sarkar --- .../ratelimiter/olric_store.go | 15 ++++++++++++- .../ratelimiter/olric_store_test.go | 21 +++++++++++++++++++ 2 files changed, 35 insertions(+), 1 deletion(-) diff --git a/src/invocation-plane-services/ratelimiter/olric_store.go b/src/invocation-plane-services/ratelimiter/olric_store.go index dcdb138752..40a66fe15c 100644 --- a/src/invocation-plane-services/ratelimiter/olric_store.go +++ b/src/invocation-plane-services/ratelimiter/olric_store.go @@ -144,12 +144,25 @@ func (store *Store) Get(ctx context.Context, key string, rate limiter.Rate) (lim } if result.TTL() > 0 { - expiration = now.Add(time.Duration(result.TTL()) * time.Millisecond) + expiration = expirationFromTTLMillis(result.TTL()) } span.SetAttributes(attribute.String("expiration", expiration.String())) return common.GetContextFromState(now, rate, expiration, int64(value)), nil } +// expirationFromTTLMillis converts the absolute Unix-millisecond expiry that +// olric.GetResponse.TTL() returns into a time.Time. GetResponse.TTL() reports +// an absolute epoch-millisecond timestamp of when the key expires - not a +// duration remaining - matching how Olric's own internal atomic-increment +// code treats the same value (time.UnixMilli(ttl) in olric-data/olric's +// internal/dmap package). Treating it as a millisecond count and adding it +// to time.Now(), as this code previously did, produced an expiration tens +// of thousands of years in the future instead of the key's real remaining +// lifetime. +func expirationFromTTLMillis(ttlMillis int64) time.Time { + return time.UnixMilli(ttlMillis) +} + // newKeyPutOptions returns the Olric put options used whenever a key is // (re)created with value 1. Folding the TTL into that same Put call - rather // than applying it in a separate follow-up write - is what closes #1571: diff --git a/src/invocation-plane-services/ratelimiter/olric_store_test.go b/src/invocation-plane-services/ratelimiter/olric_store_test.go index e7b8da120b..bda452b555 100644 --- a/src/invocation-plane-services/ratelimiter/olric_store_test.go +++ b/src/invocation-plane-services/ratelimiter/olric_store_test.go @@ -176,3 +176,24 @@ func TestGet_NewKey_SubMillisecondPeriod_TTLRoundedUp(t *testing.T) { t.Fatalf("expected a TTL option even for a sub-millisecond period, got %d options - a positive period must never be silently treated as unbounded", got) } } + +// TestExpirationFromTTLMillis_TreatsValueAsAbsoluteEpoch is the regression +// test for CodeRabbit's "Use result.TTL() as an absolute expiry" comment on +// PR #1942. olric.GetResponse.TTL() returns an absolute Unix-millisecond +// timestamp of when a key expires (confirmed against Olric's own internal +// atomic-increment code, which does time.UnixMilli(ttl) on the same value). +// The code previously treated that value as a millisecond *duration* and +// added it to time.Now(), which would push the computed expiration tens of +// thousands of years into the future instead of returning the key's real +// remaining lifetime. +func TestExpirationFromTTLMillis_TreatsValueAsAbsoluteEpoch(t *testing.T) { + want := time.Now().Add(30 * time.Second) + ttlMillis := want.UnixMilli() + + got := expirationFromTTLMillis(ttlMillis) + + diff := got.Sub(want) + if diff < -time.Millisecond || diff > time.Millisecond { + t.Fatalf("expected expiration ~%v, got %v (diff %v) - result.TTL() is an absolute epoch-ms timestamp, not a millisecond duration from now", want, got, diff) + } +}