github.com/zuoyebang/bitalosdb@v1.1.1-0.20240516111551-79a8c4d8ce20/internal/cache/lfucache/compaction_iter.go (about)

     1  // Copyright 2021 The Bitalosdb author(hustxrb@163.com) and other contributors.
     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 lfucache
    16  
    17  import (
    18  	"bytes"
    19  	"encoding/binary"
    20  	"errors"
    21  )
    22  
    23  type iterPos int8
    24  
    25  const (
    26  	iterPosCurForward       iterPos = 0
    27  	iterPosNext             iterPos = 1
    28  	iterPosPrev             iterPos = -1
    29  	iterPosCurReverse       iterPos = -2
    30  	iterPosCurForwardPaused iterPos = 2
    31  	iterPosCurReversePaused iterPos = -3
    32  )
    33  
    34  type compactionIter struct {
    35  	iter      internalIterator
    36  	err       error
    37  	key       internalKey
    38  	value     []byte
    39  	keyBuf    []byte
    40  	valid     bool
    41  	iterKey   *internalKey
    42  	iterValue []byte
    43  	skip      bool
    44  	pos       iterPos
    45  }
    46  
    47  func newCompactionIter(iter internalIterator) *compactionIter {
    48  	i := &compactionIter{
    49  		iter: iter,
    50  	}
    51  	return i
    52  }
    53  
    54  func (i *compactionIter) First() (*internalKey, []byte) {
    55  	if i.err != nil {
    56  		return nil, nil
    57  	}
    58  	i.iterKey, i.iterValue = i.iter.First()
    59  	i.pos = iterPosNext
    60  	return i.Next()
    61  }
    62  
    63  func (i *compactionIter) Next() (*internalKey, []byte) {
    64  	if i.err != nil {
    65  		return nil, nil
    66  	}
    67  
    68  	if i.pos == iterPosCurForward {
    69  		if i.skip {
    70  			i.skipInStripe()
    71  		} else {
    72  			i.nextInStripe()
    73  		}
    74  	}
    75  
    76  	i.pos = iterPosCurForward
    77  	i.valid = false
    78  	for i.iterKey != nil {
    79  		switch i.iterKey.Kind() {
    80  		case internalKeyKindDelete, internalKeyKindSet:
    81  			i.saveKey()
    82  			i.value = i.iterValue
    83  			i.valid = true
    84  			i.skip = true
    85  			return &i.key, i.value
    86  
    87  		default:
    88  			i.err = errors.New("invalid internal key kind")
    89  			i.valid = false
    90  			return nil, nil
    91  		}
    92  	}
    93  
    94  	return nil, nil
    95  }
    96  
    97  func (i *compactionIter) skipInStripe() {
    98  	i.skip = true
    99  	var change stripeChangeType
   100  	for {
   101  		change = i.nextInStripe()
   102  		if change == sameStripeNonSkippable || change == newStripe {
   103  			break
   104  		}
   105  	}
   106  
   107  	if change == newStripe {
   108  		i.skip = false
   109  	}
   110  }
   111  
   112  func (i *compactionIter) iterNext() bool {
   113  	i.iterKey, i.iterValue = i.iter.Next()
   114  	return i.iterKey != nil
   115  }
   116  
   117  type stripeChangeType int
   118  
   119  const (
   120  	newStripe stripeChangeType = iota
   121  	sameStripeSkippable
   122  	sameStripeNonSkippable
   123  )
   124  
   125  func (i *compactionIter) nextInStripe() stripeChangeType {
   126  	if !i.iterNext() {
   127  		return newStripe
   128  	}
   129  
   130  	key := i.iterKey
   131  	if bytes.Compare(i.key.UserKey, key.UserKey) != 0 {
   132  		return newStripe
   133  	}
   134  
   135  	if key.Kind() == internalKeyKindInvalid {
   136  		return sameStripeNonSkippable
   137  	}
   138  
   139  	if i.key.Kind() == internalKeyKindSet && key.Kind() == internalKeyKindSet {
   140  		valueLen := len(i.value)
   141  		iterValueLen := len(i.iterValue)
   142  
   143  		curFreq := binary.BigEndian.Uint16(i.value[valueLen-LFU_FREQ_LENGTH:])
   144  		nextFreq := binary.BigEndian.Uint16(i.iterValue[iterValueLen-LFU_FREQ_LENGTH:])
   145  
   146  		if LFU_FREQ_MAX-curFreq >= nextFreq {
   147  			curFreq += nextFreq
   148  		} else {
   149  			curFreq = LFU_FREQ_MAX
   150  		}
   151  		binary.BigEndian.PutUint16(i.value[valueLen-LFU_FREQ_LENGTH:], curFreq)
   152  		nextFreq = 0
   153  		binary.BigEndian.PutUint16(i.iterValue[iterValueLen-LFU_FREQ_LENGTH:], nextFreq)
   154  	}
   155  
   156  	return sameStripeSkippable
   157  }
   158  
   159  func (i *compactionIter) saveKey() {
   160  	i.keyBuf = append(i.keyBuf[:0], i.iterKey.UserKey...)
   161  	i.key.UserKey = i.keyBuf
   162  	i.key.Trailer = i.iterKey.Trailer
   163  }
   164  
   165  func (i *compactionIter) Key() internalKey {
   166  	return i.key
   167  }
   168  
   169  func (i *compactionIter) Value() []byte {
   170  	return i.value
   171  }
   172  
   173  func (i *compactionIter) Valid() bool {
   174  	return i.valid
   175  }
   176  
   177  func (i *compactionIter) Error() error {
   178  	return i.err
   179  }
   180  
   181  func (i *compactionIter) Close() error {
   182  	err := i.iter.Close()
   183  	if i.err == nil {
   184  		i.err = err
   185  	}
   186  
   187  	return i.err
   188  }