ronantakizawa/codereview-bench huggingface dataset:
Repository: etcd-io/etcd
File: tests/robustness/traffic/traffic.go (Go)

Code before:

const (
	DefaultWatchQPS       = 10.0
	DefaultRevisionOffset = 100
)

type Traffic interface {
	RunKeyValueLoop(ctx context.Context, param RunTrafficLoopParam)
	RunWatchLoop(ctx context.Context, param RunWatchLoopParam)
	RunCompactLoop(ctx context.Context, param RunCompactLoopParam)
	ExpectUniqueRevision() bool
}

// runWatchLoop is a helper function that implements the common watch loop pattern.
// It creates watches at a controlled rate with random revision offsets.
func runWatchLoop(ctx context.Context, p RunWatchLoopParam, cfg watchLoopConfig) {
	limiter := rate.NewLimiter(rate.Limit(p.WatchQPS), 1)

	for {
		select {
		case <-ctx.Done():
			return
		case <-p.Finish:
			return
		default:
		}

		// Wait for rate limiter before creating a new watch
		if err := limiter.Wait(ctx); err != nil {
			return
		}

		// Get current revision to watch from
		opCtx, cancel := context.WithTimeout(ctx, RequestTimeout)
		resp, err := p.Client.Get(opCtx, cfg.getKey)
		cancel()
		if err != nil {
			continue
		}
		currentRev := resp.Header.Revision

		// Generate random offset: -RevisionOffset <= offset <= RevisionOffset
		var targetRev int64
		if p.RevisionOffset > 0 {
			offset := rand.Int63n(2*p.RevisionOffset+1) - p.RevisionOffset
			targetRev = currentRev + offset
			if targetRev < 1 {
				targetRev = 1
			}
		} else {
			targetRev = currentRev + 1

Review comment (suggestion):
Not sure what this implementation is meant to follow. Does it follow kubernetes traffic or external watch? Please focus on one implementation before making it generic and remove the original code of that type of traffic.

Code after:
	Finish                             <-chan struct{}
}

type RunCompactLoopParam struct {
	Client *client.RecordingClient
	Period time.Duration
	Finish <-chan struct{}
}

type RunWatchLoopParam struct {
	Client              *client.RecordingClient
	KeyStore            *keyStore
	Finish              <-chan struct{}
	options.WatchConfig // Embedded: Interval and RevisionOffset
}

type Traffic interface {
	RunKeyValueLoop(ctx context.Context, param RunTrafficLoopParam)
	RunWatchLoop(ctx context.Context, param RunWatchLoopParam)
	RunCompactLoop(ctx context.Context, param RunCompactLoopParam)
	ExpectUniqueRevision() bool
}

// runWatchLoop is a helper function that implements the periodic watch loop pattern.
// It spawns watches at regular intervals with random revision offsets.
func runWatchLoop(ctx context.Context, p RunWatchLoopParam, cfg watchLoopConfig) {
	for {
		select {
		case <-ctx.Done():
			return
		case <-p.Finish:
			return
		case <-time.After(p.Interval):
			// Time to spawn a new watch
		}

		go func() error {
			resp, err := p.Client.Get(ctx, cfg.watchKey)
			if err != nil {
				return err
			}
			rev := resp.Header.Revision + p.RevisionOffset

			watchCtx, cancel := context.WithCancel(ctx)
			defer cancel()
			w := p.Client.Watch(watchCtx, cfg.watchKey, rev, true, true, true)
			for {
				select {
				case <-ctx.Done():
					return ctx.Err()
				case <-p.Finish:
Repository: etcd-io/etcd
File: tests/robustness/traffic/kubernetes.go (Go)

Code before:
	g.Wait()
}

func (t kubernetesTraffic) RunWatchLoop(ctx context.Context, p RunWatchLoopParam) {
	s := p.Storage
	keyPrefix := "/registry/" + t.resource + "/"
	for {
		select {
		case <-ctx.Done():
			return
		case <-p.Finish:
			return
		case <-time.After(DefaultWatchInterval):
			// Time to spawn a new watch
		}

		p.WaitGroup.Add(1)
		go func() {
			defer p.WaitGroup.Done()
			resp, err := p.Client.Get(ctx, keyPrefix)
			if err != nil {
				return
			}
			rev := resp.Header.Revision + DefaultRevisionOffset

			watchCtx, cancel := context.WithCancel(ctx)
			defer cancel()

			// Kubernetes issues Watch requests by requiring a leader to exist
			watchCtx = clientv3.WithRequireLeader(watchCtx)
			w := p.Client.Watch(watchCtx, keyPrefix, rev, true, true, true)
			for {
				select {
				case <-ctx.Done():
					return
				case <-p.Finish:
					return
				case resp, ok := <-w:
					if !ok {
						return
					}
					s.Update(resp)
				}
			}
		}()
	}
}

func (t kubernetesTraffic) Read(ctx context.Context, c *client.RecordingClient, s *storage, limiter *rate.Limiter, keyPrefix string) error {
	kc := kubernetes.Client{Client: &clientv3.Client{KV: c}}
	op := random.PickRandom(t.readChoices)

Review comment (suggestion):
```suggestion
			watchCtx, cancel := context.WithTimeout(ctx, WatchTimeout)
```

Code after:
	g.Wait()
}

func (t kubernetesTraffic) RunWatchLoop(ctx context.Context, p RunWatchLoopParam) {
	s := p.Storage
	keyPrefix := "/registry/" + t.resource + "/"
	for {
		select {
		case <-ctx.Done():
			return
		case <-p.Finish:
			return
		case <-time.After(DefaultWatchInterval):
			// Time to spawn a new watch
		}

		p.WaitGroup.Add(1)
		go func() {
			defer p.WaitGroup.Done()
			resp, err := p.Client.Get(ctx, keyPrefix)
			if err != nil {
				return
			}
			rev := resp.Header.Revision + DefaultRevisionOffset

			watchCtx, cancel := context.WithTimeout(ctx, WatchTimeout)
			defer cancel()

			// Kubernetes issues Watch requests by requiring a leader to exist
			watchCtx = clientv3.WithRequireLeader(watchCtx)
			w := p.Client.Watch(watchCtx, keyPrefix, rev, true, true, true)
			for {
				select {
				case <-ctx.Done():
					return
				case <-p.Finish:
					return
				case resp, ok := <-w:
					if !ok {
						return
					}
					s.Update(resp)
				}
			}
		}()
	}
}

func (t kubernetesTraffic) Read(ctx context.Context, c *client.RecordingClient, s *storage, limiter *rate.Limiter, keyPrefix string) error {
	kc := kubernetes.Client{Client: &clientv3.Client{KV: c}}
	op := random.PickRandom(t.readChoices)
Repository: etcd-io/etcd
File: etcdutl/etcdutl/hashkv_command_test.go (Go)

Code before:
// Copyright 2024 The etcd Authors
//
// 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 etcdutl

import (
	"path/filepath"
	"testing"

	"go.uber.org/zap"
	"gotest.tools/v3/assert"

	"go.etcd.io/etcd/server/v3/lease"
	"go.etcd.io/etcd/server/v3/storage/backend"
	"go.etcd.io/etcd/server/v3/storage/mvcc"

Review comment (security):
```suggestion
// Copyright 2025 The etcd Authors
```

Code after:
// Copyright 2026 The etcd Authors
//
// 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 etcdutl

import (
	"path/filepath"
	"testing"

	"go.uber.org/zap"
	"gotest.tools/v3/assert"

	"go.etcd.io/etcd/server/v3/lease"
	"go.etcd.io/etcd/server/v3/storage/backend"
	"go.etcd.io/etcd/server/v3/storage/mvcc"
Repository: etcd-io/etcd
File: etcdutl/etcdutl/hashkv_command_test.go (Go)

Code before:
func TestCalculateHashKV_WithData(t *testing.T) {
	tempDir := t.TempDir()
	dbPath := filepath.Join(tempDir, "test_with_data.db")

	// Create database with mvcc store
	b := backend.NewDefaultBackend(zap.NewNop(), dbPath)
	st := mvcc.NewStore(zap.NewNop(), b, &lease.FakeLessor{}, mvcc.StoreConfig{})

	// Put some test data
	st.Put([]byte("test-key"), []byte("test-value"), 1)

	b.Close()

	// Test with revision 0 (latest)
	result, err := calculateHashKV(dbPath, 0)
	if err != nil {
		t.Fatalf("calculateHashKV() failed: %v", err)
	}

	assert.Equal(t, uint32(645561629), result.Hash)
	assert.Equal(t, int64(2), result.HashRevision)
	assert.Equal(t, int64(-1), result.CompactRevision)

	t.Logf("With data - Hash: %d, HashRevision: %d, CompactRevision: %d",
		result.Hash, result.HashRevision, result.CompactRevision)
}

func TestCalculateHashKV_InvalidRevision(t *testing.T) {
	tempDir := t.TempDir()
	dbPath := filepath.Join(tempDir, "test_invalid_rev.db")

	b := backend.NewDefaultBackend(zap.NewNop(), dbPath)
	st := mvcc.NewStore(zap.NewNop(), b, &lease.FakeLessor{}, mvcc.StoreConfig{})
	st.Put([]byte("key"), []byte("value"), 1)
	b.Close()

	// Test with non-existent revision
	_, err := calculateHashKV(dbPath, 999)
	assert.ErrorContains(t, err, "required revision is a future revision")
}


Review comment (suggestion):
These two cases should be able to be merged into one case, which contains a table-driven case (2 sub test)

Code after:
				st := mvcc.NewStore(zap.NewNop(), b, &lease.FakeLessor{}, mvcc.StoreConfig{})
				st.Put([]byte("test-key"), []byte("test-value"), 1)
				st.Close()
				b.Close()

				return dbPath, func() {}
			},
			revision:        0,
			expectedHash:    645561629,
			expectedRev:     2,
			expectedCompact: -1,
			expectError:     false,
		},
		{
			name: "invalid revision",
			setupFunc: func(t *testing.T) (string, func()) {
				tempDir := t.TempDir()
				dbPath := filepath.Join(tempDir, "test_invalid_rev.db")

				b := backend.NewDefaultBackend(zap.NewNop(), dbPath)
				st := mvcc.NewStore(zap.NewNop(), b, &lease.FakeLessor{}, mvcc.StoreConfig{})
				st.Put([]byte("key"), []byte("value"), 1)
				st.Close()
				b.Close()

				return dbPath, func() {}
			},
			revision:      999,
			expectError:   true,
			errorContains: "required revision is a future revision",
		},
	}

	for _, tc := range testCases {
		t.Run(tc.name, func(t *testing.T) {
			defer func() {
				if r := recover(); r != nil {
					t.Logf("Recovered from panic: %v", r)
				}
			}()

			dbPath, cleanup := tc.setupFunc(t)
			defer cleanup()

			result, err := calculateHashKV(dbPath, tc.revision)

			if tc.expectError {
				assert.Assert(t, err != nil)
				if tc.errorContains != "" {
					assert.ErrorContains(t, err, tc.errorContains)
				}
Repository: etcd-io/etcd
File: scripts/fix/mod-tidy.sh (Shell)

Code before:
#!/usr/bin/env bash

# Copyright 2024 The etcd Authors
#
# 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.

set -euo pipefail

source ./scripts/test_lib.sh

log_callout "Tidying go.mod files"
run_for_all_workspace_modules run go mod tidy
log_success "go.mod files tidied"


Review comment (security):
isn't this supposed to be ``` 2026 The etcd Authors```?

Code after:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)}
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh"

log_callout "Tidying go.mod files"
run_for_workspace_modules run sh -c "rm -f ./go.sum && go mod tidy"
log_success "go.mod files tidied"
Repository: etcd-io/etcd
File: scripts/fix/mod-tidy.sh (Shell)

Code before:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)} 
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh" 

log_callout "Tidying go.mod files"
run_for_workspace_modules go mod tidy
log_success "go.mod files tidied"


Review comment (suggestion):
Please refer to how we're currently running the go mod tidy. Note that we're removing the `go.mod` file. I'm unsure why we're doing that.

https://github.com/etcd-io/etcd/blob/566a42c031875a5d4ef9197f0fb197ce98c6afb6/scripts/test.sh#L629-L636

I believe it would also be a good idea to delete those two functions from the `scripts/test.sh` file, so we start cleaning up that big file.



Code after:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)}
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh"

log_callout "Tidying go.mod files"
run_for_workspace_modules run sh -c "rm -f ./go.sum && go mod tidy"
log_success "go.mod files tidied"
Repository: etcd-io/etcd
File: scripts/fix/mod-tidy.sh (Shell)

Code before:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)} 
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh" 

log_callout "Tidying go.mod files"
run_for_workspace_modules sh -c "rm -f ./go.sum && go mod tidy"
log_success "go.mod files tidied"


Review comment (suggestion):
Can you please remove the trailing whitespace from these two lines? It only became apparent after checking out the code locally.

```suggestion
ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)}
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh"
```

Code after:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)}
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh"

log_callout "Tidying go.mod files"
run_for_workspace_modules run sh -c "rm -f ./go.sum && go mod tidy"
log_success "go.mod files tidied"
Repository: etcd-io/etcd
File: scripts/fix/mod-tidy.sh (Shell)

Code before:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)} 
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh" 

log_callout "Tidying go.mod files"
run_for_workspace_modules sh -c "rm -f ./go.sum && go mod tidy"
log_success "go.mod files tidied"


Review comment (suggestion):
Let's change this line to wrap it with the `run` function, so we add verbosity on what the command is doing.

```suggestion
run_for_workspace_modules run sh -c "rm -f ./go.sum && go mod tidy"
```

Code after:
#!/usr/bin/env bash

# Copyright 2026 The etcd Authors
#
# 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.

set -euo pipefail

ETCD_ROOT_DIR=${ETCD_ROOT_DIR:-$(git rev-parse --show-toplevel)}
source "${ETCD_ROOT_DIR}/scripts/test_lib.sh"

log_callout "Tidying go.mod files"
run_for_workspace_modules run sh -c "rm -f ./go.sum && go mod tidy"
log_success "go.mod files tidied"
Repository: etcd-io/etcd
File: server/etcdserver/v3_server.go (Go)

Code before:
	for cctx.Err() == nil {
		leader, lerr := s.waitLeader(cctx)
		if lerr != nil {
			return -1, lerr
		}
		for _, url := range leader.PeerURLs {
			lurl := url + leasehttp.LeasePrefix
			ttl, err := leasehttp.RenewHTTP(cctx, id, lurl, s.peerRt)
			if err == nil || errorspkg.Is(err, lease.ErrLeaseNotFound) {
				return ttl, err
			}
		}
		// Throttle in case of e.g. connection problems.
		select {
		case <-time.After(50 * time.Millisecond):
		case <-cctx.Done():
		}
	}

	err := cctx.Err()
	switch {
	case errorspkg.Is(err, context.DeadlineExceeded):
		return -1, errors.ErrTimeout
	case errorspkg.Is(err, context.Canceled):
		return -1, errors.ErrCanceled
	// Should be unreachable, but we keep it defensive.
	default:
		s.Logger().Warn("Unexpected lease renew context error", zap.Error(err))
		return -1, errors.ErrTimeout
	}
}

func (s *EtcdServer) checkLeaseTimeToLive(ctx context.Context, leaseID lease.LeaseID) (uint64, error) {
	rev := s.AuthStore().Revision()
	if !s.AuthStore().IsAuthEnabled() {
		return rev, nil
	}
	authInfo, err := s.AuthInfoFromCtx(ctx)
	if err != nil {
		return rev, err
	}
	if authInfo == nil {
		return rev, auth.ErrUserEmpty
	}

	l := s.lessor.Lookup(leaseID)
	if l != nil {
		for _, key := range l.Keys() {
			if err := s.AuthStore().IsRangePermitted(authInfo, []byte(key), []byte{}); err != nil {
				return 0, err
			}

Review comment (suggestion):
The log "Unexpected lease renew ..." already conveys this comment.
```suggestion
```

Code after:
	for cctx.Err() == nil {
		leader, lerr := s.waitLeader(cctx)
		if lerr != nil {
			return -1, lerr
		}
		for _, url := range leader.PeerURLs {
			lurl := url + leasehttp.LeasePrefix
			ttl, err := leasehttp.RenewHTTP(cctx, id, lurl, s.peerRt)
			if err == nil || errorspkg.Is(err, lease.ErrLeaseNotFound) {
				return ttl, err
			}
		}
		// Throttle in case of e.g. connection problems.
		time.Sleep(50 * time.Millisecond)
	}

	err := cctx.Err()
	switch {
	case errorspkg.Is(err, context.DeadlineExceeded):
		return -1, errors.ErrTimeout
	case errorspkg.Is(err, context.Canceled):
		return -1, errors.ErrCanceled
	default:
		s.Logger().Warn("Unexpected lease renew context error", zap.Error(err))
		return -1, errors.ErrCanceled
	}
}

func (s *EtcdServer) checkLeaseTimeToLive(ctx context.Context, leaseID lease.LeaseID) (uint64, error) {
	rev := s.AuthStore().Revision()
	if !s.AuthStore().IsAuthEnabled() {
		return rev, nil
	}
	authInfo, err := s.AuthInfoFromCtx(ctx)
	if err != nil {
		return rev, err
	}
	if authInfo == nil {
		return rev, auth.ErrUserEmpty
	}

	l := s.lessor.Lookup(leaseID)
	if l != nil {
		for _, key := range l.Keys() {
			if err := s.AuthStore().IsRangePermitted(authInfo, []byte(key), []byte{}); err != nil {
				return 0, err
			}
		}
	}

	return rev, nil
Repository: etcd-io/etcd
File: cache/snapshot.go (Go)

Code before:
// See the License for the specific language governing permissions and
// limitations under the License.

package cache

import (
	"github.com/google/btree"

	"go.etcd.io/etcd/api/v3/mvccpb"
)

// snapshot captures a full, point-in-time view of the KV state at rev.
type snapshot struct {
	rev  int64
	tree *btree.BTree
}

func newClonedSnapshot(rev int64, t *btree.BTree) *snapshot {
	return &snapshot{rev: rev, tree: t.Clone()}
}

func (s *snapshot) Range(startKey, endKey []byte) []*mvccpb.KeyValue {
	var out []*mvccpb.KeyValue
	switch {
	case len(endKey) == 0:
		if item := s.tree.Get(probeItemFromKey(startKey)); item != nil {
			out = append(out, item.(*kvItem).kv)
		}
	case isPrefixScan(endKey):
		s.tree.AscendGreaterOrEqual(probeItemFromKey(startKey), func(item btree.Item) bool {
			out = append(out, item.(*kvItem).kv)
			return true
		})
	default:
		s.tree.AscendRange(probeItemFromKey(startKey), probeItemFromKey(endKey), func(item btree.Item) bool {
			out = append(out, item.(*kvItem).kv)
			return true
		})
	}
	return out
}

func isPrefixScan(endKey []byte) bool {
	return len(endKey) == 1 && endKey[0] == 0
}

func probeItemFromKey(key []byte) *kvItem { return &kvItem{key: string(key)} }


Review comment (suggestion):
nit
```suggestion
		if item := s.tree.Get(probeItemFromKey(startKey)); item != nil {
```

Code after:
// See the License for the specific language governing permissions and
// limitations under the License.

package cache

import (
	"k8s.io/utils/third_party/forked/golang/btree"

	"go.etcd.io/etcd/api/v3/mvccpb"
)

// snapshot captures a full, point-in-time view of the KV state at rev.
type snapshot struct {
	rev  int64
	tree *btree.BTree[*kvItem]
}

func newClonedSnapshot(rev int64, t *btree.BTree[*kvItem]) *snapshot {
	return &snapshot{rev: rev, tree: t.Clone()}
}

func (s *snapshot) Range(startKey, endKey []byte) []*mvccpb.KeyValue {
	var out []*mvccpb.KeyValue
	switch {
	case len(endKey) == 0:
		if item, ok := s.tree.Get(probeItemFromKey(startKey)); ok {
			out = append(out, item.kv)
		}
	case isPrefixScan(endKey):
		s.tree.AscendGreaterOrEqual(probeItemFromKey(startKey), func(item *kvItem) bool {
			out = append(out, item.kv)
			return true
		})
	default:
		s.tree.AscendRange(
			probeItemFromKey(startKey),
			probeItemFromKey(endKey),
			func(item *kvItem) bool {
				out = append(out, item.kv)
				return true
			},
		)
	}
	return out
}

func isPrefixScan(endKey []byte) bool {
	return len(endKey) == 1 && endKey[0] == 0
}

func probeItemFromKey(key []byte) *kvItem { return &kvItem{key: string(key)} }