github.com/whtcorpsinc/MilevaDB-Prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/ekv/utils.go (about)

     1  // Copyright 2020 WHTCORPS INC, Inc.
     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  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package ekv
    15  
    16  import (
    17  	"context"
    18  	"strconv"
    19  
    20  	"github.com/whtcorpsinc/errors"
    21  )
    22  
    23  // IncInt64 increases the value for key k in ekv causetstore by step.
    24  func IncInt64(rm RetrieverMutator, k Key, step int64) (int64, error) {
    25  	val, err := rm.Get(context.TODO(), k)
    26  	if IsErrNotFound(err) {
    27  		err = rm.Set(k, []byte(strconv.FormatInt(step, 10)))
    28  		if err != nil {
    29  			return 0, err
    30  		}
    31  		return step, nil
    32  	}
    33  	if err != nil {
    34  		return 0, err
    35  	}
    36  
    37  	intVal, err := strconv.ParseInt(string(val), 10, 64)
    38  	if err != nil {
    39  		return 0, errors.Trace(err)
    40  	}
    41  
    42  	intVal += step
    43  	err = rm.Set(k, []byte(strconv.FormatInt(intVal, 10)))
    44  	if err != nil {
    45  		return 0, err
    46  	}
    47  	return intVal, nil
    48  }
    49  
    50  // GetInt64 get int64 value which created by IncInt64 method.
    51  func GetInt64(ctx context.Context, r Retriever, k Key) (int64, error) {
    52  	val, err := r.Get(ctx, k)
    53  	if IsErrNotFound(err) {
    54  		return 0, nil
    55  	}
    56  	if err != nil {
    57  		return 0, err
    58  	}
    59  	intVal, err := strconv.ParseInt(string(val), 10, 64)
    60  	if err != nil {
    61  		return intVal, errors.Trace(err)
    62  	}
    63  	return intVal, nil
    64  }
    65  
    66  // WalkMemBuffer iterates all buffered ekv pairs in memBuf
    67  func WalkMemBuffer(memBuf Retriever, f func(k Key, v []byte) error) error {
    68  	iter, err := memBuf.Iter(nil, nil)
    69  	if err != nil {
    70  		return errors.Trace(err)
    71  	}
    72  
    73  	defer iter.Close()
    74  	for iter.Valid() {
    75  		if err = f(iter.Key(), iter.Value()); err != nil {
    76  			return errors.Trace(err)
    77  		}
    78  		err = iter.Next()
    79  		if err != nil {
    80  			return errors.Trace(err)
    81  		}
    82  	}
    83  
    84  	return nil
    85  }
    86  
    87  // BufferBatchGetter is the type for BatchGet with MemBuffer.
    88  type BufferBatchGetter struct {
    89  	buffer   MemBuffer
    90  	midbse   Getter
    91  	snapshot Snapshot
    92  }
    93  
    94  // NewBufferBatchGetter creates a new BufferBatchGetter.
    95  func NewBufferBatchGetter(buffer MemBuffer, midbseCache Getter, snapshot Snapshot) *BufferBatchGetter {
    96  	return &BufferBatchGetter{buffer: buffer, midbse: midbseCache, snapshot: snapshot}
    97  }
    98  
    99  // BatchGet implements the BatchGetter interface.
   100  func (b *BufferBatchGetter) BatchGet(ctx context.Context, keys []Key) (map[string][]byte, error) {
   101  	if b.buffer.Len() == 0 {
   102  		return b.snapshot.BatchGet(ctx, keys)
   103  	}
   104  	bufferValues := make([][]byte, len(keys))
   105  	shrinkKeys := make([]Key, 0, len(keys))
   106  	for i, key := range keys {
   107  		val, err := b.buffer.Get(ctx, key)
   108  		if err == nil {
   109  			bufferValues[i] = val
   110  			continue
   111  		}
   112  		if !IsErrNotFound(err) {
   113  			return nil, errors.Trace(err)
   114  		}
   115  		if b.midbse != nil {
   116  			val, err = b.midbse.Get(ctx, key)
   117  			if err == nil {
   118  				bufferValues[i] = val
   119  				continue
   120  			}
   121  		}
   122  		shrinkKeys = append(shrinkKeys, key)
   123  	}
   124  	storageValues, err := b.snapshot.BatchGet(ctx, shrinkKeys)
   125  	if err != nil {
   126  		return nil, errors.Trace(err)
   127  	}
   128  	for i, key := range keys {
   129  		if len(bufferValues[i]) == 0 {
   130  			continue
   131  		}
   132  		storageValues[string(key)] = bufferValues[i]
   133  	}
   134  	return storageValues, nil
   135  }