github.com/whtcorpsinc/milevadb-prod@v0.0.0-20211104133533-f57f4be3b597/causetstore/causetstore/entangledstore/raw_handler.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 entangledstore
    15  
    16  import (
    17  	"bytes"
    18  	"context"
    19  	"sync"
    20  
    21  	"github.com/ngaut/entangledstore/lockstore"
    22  	"github.com/whtcorpsinc/ekvproto/pkg/ekvrpcpb"
    23  )
    24  
    25  type rawHandler struct {
    26  	mu          sync.RWMutex
    27  	causetstore *lockstore.MemStore
    28  }
    29  
    30  func newRawHandler() *rawHandler {
    31  	return &rawHandler{
    32  		causetstore: lockstore.NewMemStore(4096),
    33  	}
    34  }
    35  
    36  func (h *rawHandler) RawGet(_ context.Context, req *ekvrpcpb.RawGetRequest) (*ekvrpcpb.RawGetResponse, error) {
    37  	h.mu.RLock()
    38  	defer h.mu.RUnlock()
    39  	val := h.causetstore.Get(req.Key, nil)
    40  	return &ekvrpcpb.RawGetResponse{
    41  		Value:    val,
    42  		NotFound: len(val) == 0,
    43  	}, nil
    44  }
    45  
    46  func (h *rawHandler) RawBatchGet(_ context.Context, req *ekvrpcpb.RawBatchGetRequest) (*ekvrpcpb.RawBatchGetResponse, error) {
    47  	h.mu.RLock()
    48  	defer h.mu.RUnlock()
    49  	pairs := make([]*ekvrpcpb.EkvPair, len(req.Keys))
    50  	for i, key := range req.Keys {
    51  		pairs[i] = &ekvrpcpb.EkvPair{
    52  			Key:   key,
    53  			Value: h.causetstore.Get(key, nil),
    54  		}
    55  	}
    56  	return &ekvrpcpb.RawBatchGetResponse{Pairs: pairs}, nil
    57  }
    58  
    59  func (h *rawHandler) RawPut(_ context.Context, req *ekvrpcpb.RawPutRequest) (*ekvrpcpb.RawPutResponse, error) {
    60  	h.mu.Lock()
    61  	defer h.mu.Unlock()
    62  	h.causetstore.Put(req.Key, req.Value)
    63  	return &ekvrpcpb.RawPutResponse{}, nil
    64  }
    65  
    66  func (h *rawHandler) RawBatchPut(_ context.Context, req *ekvrpcpb.RawBatchPutRequest) (*ekvrpcpb.RawBatchPutResponse, error) {
    67  	h.mu.Lock()
    68  	defer h.mu.Unlock()
    69  	for _, pair := range req.Pairs {
    70  		h.causetstore.Put(pair.Key, pair.Value)
    71  	}
    72  	return &ekvrpcpb.RawBatchPutResponse{}, nil
    73  }
    74  
    75  func (h *rawHandler) RawDelete(_ context.Context, req *ekvrpcpb.RawDeleteRequest) (*ekvrpcpb.RawDeleteResponse, error) {
    76  	h.mu.Lock()
    77  	defer h.mu.Unlock()
    78  	h.causetstore.Delete(req.Key)
    79  	return &ekvrpcpb.RawDeleteResponse{}, nil
    80  }
    81  
    82  func (h *rawHandler) RawBatchDelete(_ context.Context, req *ekvrpcpb.RawBatchDeleteRequest) (*ekvrpcpb.RawBatchDeleteResponse, error) {
    83  	h.mu.Lock()
    84  	defer h.mu.Unlock()
    85  	for _, key := range req.Keys {
    86  		h.causetstore.Delete(key)
    87  	}
    88  	return &ekvrpcpb.RawBatchDeleteResponse{}, nil
    89  }
    90  
    91  func (h *rawHandler) RawDeleteRange(_ context.Context, req *ekvrpcpb.RawDeleteRangeRequest) (*ekvrpcpb.RawDeleteRangeResponse, error) {
    92  	h.mu.Lock()
    93  	defer h.mu.Unlock()
    94  	it := h.causetstore.NewIterator()
    95  	var keys [][]byte
    96  	for it.Seek(req.StartKey); it.Valid(); it.Next() {
    97  		if bytes.Compare(it.Key(), req.EndKey) >= 0 {
    98  			break
    99  		}
   100  		keys = append(keys, safeCopy(it.Key()))
   101  	}
   102  	for _, key := range keys {
   103  		h.causetstore.Delete(key)
   104  	}
   105  	return &ekvrpcpb.RawDeleteRangeResponse{}, nil
   106  }
   107  
   108  func (h *rawHandler) RawScan(_ context.Context, req *ekvrpcpb.RawScanRequest) (*ekvrpcpb.RawScanResponse, error) {
   109  	h.mu.RLock()
   110  	defer h.mu.RUnlock()
   111  	it := h.causetstore.NewIterator()
   112  	var pairs []*ekvrpcpb.EkvPair
   113  	if !req.Reverse {
   114  		for it.Seek(req.StartKey); it.Valid(); it.Next() {
   115  			if len(pairs) >= int(req.Limit) {
   116  				break
   117  			}
   118  			if len(req.EndKey) > 0 && bytes.Compare(it.Key(), req.EndKey) >= 0 {
   119  				break
   120  			}
   121  			pairs = h.appendPair(pairs, it)
   122  		}
   123  	} else {
   124  		for it.SeekForPrev(req.StartKey); it.Valid(); it.Prev() {
   125  			if bytes.Equal(it.Key(), req.StartKey) {
   126  				continue
   127  			}
   128  			if len(pairs) >= int(req.Limit) {
   129  				break
   130  			}
   131  			if bytes.Compare(it.Key(), req.EndKey) < 0 {
   132  				break
   133  			}
   134  			pairs = h.appendPair(pairs, it)
   135  		}
   136  	}
   137  	return &ekvrpcpb.RawScanResponse{Ekvs: pairs}, nil
   138  }
   139  
   140  func (h *rawHandler) appendPair(pairs []*ekvrpcpb.EkvPair, it *lockstore.Iterator) []*ekvrpcpb.EkvPair {
   141  	pair := &ekvrpcpb.EkvPair{
   142  		Key:   safeCopy(it.Key()),
   143  		Value: safeCopy(it.Value()),
   144  	}
   145  	return append(pairs, pair)
   146  }
   147  
   148  func safeCopy(val []byte) []byte {
   149  	return append([]byte{}, val...)
   150  }