khenaidoo | d948f77 | 2021-08-11 17:49:24 -0400 | [diff] [blame] | 1 | // Copyright 2017 The etcd Authors |
| 2 | // |
| 3 | // Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | // you may not use this file except in compliance with the License. |
| 5 | // You may obtain a copy of the License at |
| 6 | // |
| 7 | // http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | // |
| 9 | // Unless required by applicable law or agreed to in writing, software |
| 10 | // distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | // See the License for the specific language governing permissions and |
| 13 | // limitations under the License. |
| 14 | |
| 15 | package backend |
| 16 | |
| 17 | import ( |
| 18 | "bytes" |
| 19 | "math" |
| 20 | "sync" |
| 21 | |
| 22 | bolt "github.com/coreos/bbolt" |
| 23 | ) |
| 24 | |
| 25 | // safeRangeBucket is a hack to avoid inadvertently reading duplicate keys; |
| 26 | // overwrites on a bucket should only fetch with limit=1, but safeRangeBucket |
| 27 | // is known to never overwrite any key so range is safe. |
| 28 | var safeRangeBucket = []byte("key") |
| 29 | |
| 30 | type ReadTx interface { |
| 31 | Lock() |
| 32 | Unlock() |
| 33 | |
| 34 | UnsafeRange(bucketName []byte, key, endKey []byte, limit int64) (keys [][]byte, vals [][]byte) |
| 35 | UnsafeForEach(bucketName []byte, visitor func(k, v []byte) error) error |
| 36 | } |
| 37 | |
| 38 | type readTx struct { |
| 39 | // mu protects accesses to the txReadBuffer |
| 40 | mu sync.RWMutex |
| 41 | buf txReadBuffer |
| 42 | |
| 43 | // txmu protects accesses to buckets and tx on Range requests. |
| 44 | txmu sync.RWMutex |
| 45 | tx *bolt.Tx |
| 46 | buckets map[string]*bolt.Bucket |
| 47 | } |
| 48 | |
| 49 | func (rt *readTx) Lock() { rt.mu.RLock() } |
| 50 | func (rt *readTx) Unlock() { rt.mu.RUnlock() } |
| 51 | |
| 52 | func (rt *readTx) UnsafeRange(bucketName, key, endKey []byte, limit int64) ([][]byte, [][]byte) { |
| 53 | if endKey == nil { |
| 54 | // forbid duplicates for single keys |
| 55 | limit = 1 |
| 56 | } |
| 57 | if limit <= 0 { |
| 58 | limit = math.MaxInt64 |
| 59 | } |
| 60 | if limit > 1 && !bytes.Equal(bucketName, safeRangeBucket) { |
| 61 | panic("do not use unsafeRange on non-keys bucket") |
| 62 | } |
| 63 | keys, vals := rt.buf.Range(bucketName, key, endKey, limit) |
| 64 | if int64(len(keys)) == limit { |
| 65 | return keys, vals |
| 66 | } |
| 67 | |
| 68 | // find/cache bucket |
| 69 | bn := string(bucketName) |
| 70 | rt.txmu.RLock() |
| 71 | bucket, ok := rt.buckets[bn] |
| 72 | rt.txmu.RUnlock() |
| 73 | if !ok { |
| 74 | rt.txmu.Lock() |
| 75 | bucket = rt.tx.Bucket(bucketName) |
| 76 | rt.buckets[bn] = bucket |
| 77 | rt.txmu.Unlock() |
| 78 | } |
| 79 | |
| 80 | // ignore missing bucket since may have been created in this batch |
| 81 | if bucket == nil { |
| 82 | return keys, vals |
| 83 | } |
| 84 | rt.txmu.Lock() |
| 85 | c := bucket.Cursor() |
| 86 | rt.txmu.Unlock() |
| 87 | |
| 88 | k2, v2 := unsafeRange(c, key, endKey, limit-int64(len(keys))) |
| 89 | return append(k2, keys...), append(v2, vals...) |
| 90 | } |
| 91 | |
| 92 | func (rt *readTx) UnsafeForEach(bucketName []byte, visitor func(k, v []byte) error) error { |
| 93 | dups := make(map[string]struct{}) |
| 94 | getDups := func(k, v []byte) error { |
| 95 | dups[string(k)] = struct{}{} |
| 96 | return nil |
| 97 | } |
| 98 | visitNoDup := func(k, v []byte) error { |
| 99 | if _, ok := dups[string(k)]; ok { |
| 100 | return nil |
| 101 | } |
| 102 | return visitor(k, v) |
| 103 | } |
| 104 | if err := rt.buf.ForEach(bucketName, getDups); err != nil { |
| 105 | return err |
| 106 | } |
| 107 | rt.txmu.Lock() |
| 108 | err := unsafeForEach(rt.tx, bucketName, visitNoDup) |
| 109 | rt.txmu.Unlock() |
| 110 | if err != nil { |
| 111 | return err |
| 112 | } |
| 113 | return rt.buf.ForEach(bucketName, visitor) |
| 114 | } |
| 115 | |
| 116 | func (rt *readTx) reset() { |
| 117 | rt.buf.reset() |
| 118 | rt.buckets = make(map[string]*bolt.Bucket) |
| 119 | rt.tx = nil |
| 120 | } |