Replace mem KV with Badger in-memory KV #9790
Conversation
github-actions
bot
added
area/gateway
arielshaqed
added
minor-change
arielshaqed
requested review from
Copilot and
nopcoder
and removed request for
nopcoder
There was a problem hiding this comment.
Pull request overview
This PR replaces a custom in-memory key-value store implementation with Badger's built-in in-memory mode. The change consolidates the codebase by using the same underlying storage engine (Badger) for both local file-based and in-memory KV stores.
Key changes:
- Removed the custom
pkg/kv/mempackage (store.go and store_test.go) - Added
MemDrivertopkg/kv/local/driver.goalongside the existingLocalDriver - Updated all imports across the codebase from
pkg/kv/memtopkg/kv/local
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| pkg/kv/mem/store.go | Deleted custom in-memory KV implementation using sortedmap |
| pkg/kv/mem/store_test.go | Deleted tests for custom mem KV implementation |
| pkg/kv/local/driver.go | Added MemDriver that creates Badger in-memory stores; renamed DriverName constants |
| pkg/kv/local/store.go | Updated mutex name from driverLock to localDriverLock |
| pkg/kv/local/store_test.go | Added TestMemKV function; updated driver name references |
| pkg/kv/msg_test.go | Updated driver name references and added benchmark test |
| pkg/kv/kvtest/store.go | Updated import from mem to local package |
| pkg/stats/usage_counter_test.go | Removed unused mem import |
| pkg/gateway/testutil/gateway_setup.go | Updated import from mem to local package |
| pkg/gateway/multipart/main_test.go | Removed unused mem import |
| pkg/distributed/mc_owner_test.go | Updated import from mem to local package |
| pkg/api/serve_test.go | Updated driver name reference and import |
| pkg/actions/service_test.go | Removed unused mem import |
| contrib/auth/acl/cmd/acl/cmd/run.go | Removed mem import |
| cmd/lakefs/cmd/run.go | Updated driver name references and removed mem import |
| cmd/lakefs/cmd/root.go | Updated driver name references and removed mem import |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| Type: local.MemDriverName, | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("failed to open kv '%s' store: %s", local.LocalDriverName, err) |
There was a problem hiding this comment.
The error message uses the wrong driver name. It should use local.MemDriverName instead of local.LocalDriverName to match the driver type being tested.
| t.Fatalf("failed to open kv '%s' store: %s", local.LocalDriverName, err) | |
| t.Fatalf("failed to open kv '%s' store: %s", local.MemDriverName, err) |
| b.Run("mem (run with --benchmem)", func(b *testing.B) { | ||
| source := rand.NewSource(randomGetSeed) | ||
| b.ResetTimer() | ||
| testGetMsgs(b, ctx, store, b.N, int32(b.N), source) | ||
| }) |
There was a problem hiding this comment.
This benchmark test is named "mem (run with --benchmem)" but it appears to be running within the context of all driver benchmarks (mem, local, postgres, dynamoDB). The name is confusing because it suggests it's specifically for the mem driver, but it's actually nested inside each driver's benchmark loop. This benchmark should either be renamed to clarify what it's testing, or moved to the appropriate location if it's only meant for the mem driver.
| func (d *MemDriver) Open(ctx context.Context, _ kvparams.Config) (kv.Store, error) { | ||
| // no database open for this path | ||
| logger := logging.FromContext(ctx).WithField("store", "mem") | ||
| opts := badger.DefaultOptions("").WithInMemory(true) | ||
| opts.Logger = &BadgerLogger{logger} | ||
| db, err := badger.Open(opts) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
| connection := &Store{ | ||
| db: db, | ||
| logger: logger, | ||
| prefetchSize: memPrefetchSize, | ||
| } | ||
| connection.refCount++ | ||
| return connection, nil | ||
| } |
There was a problem hiding this comment.
The MemDriver creates a Store without setting the path field, which causes a problem in the Store.Close() method. When Close() is called, it attempts to delete from dbMap using s.path as the key, but since path is empty for mem stores, this will try to delete the empty string key from dbMap. Additionally, mem stores are never added to dbMap, so they shouldn't be trying to remove themselves from it. The refCount management also doesn't make sense for mem stores since each call to Open creates a new independent in-memory database. Consider either setting a unique path identifier for mem stores or having MemDriver create a different Store type that doesn't use dbMap-based lifecycle management.
|
Blocked by failing tests - #9793. |
Raised by @nopcoder in #9766.
Rather than use a home-brewed (but fast!) KV, use Badger - same as "local" - with in-memory mode.