github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/kv/kvserver/closedts/provider/testutils/storage.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package testutils
    12  
    13  import (
    14  	"sort"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/closedts/ctpb"
    17  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    18  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    19  )
    20  
    21  // TestStorage is a storage backed by a map[NodeID]Entry.
    22  type TestStorage struct {
    23  	mu syncutil.Mutex
    24  	m  map[roachpb.NodeID][]ctpb.Entry
    25  }
    26  
    27  // VisitAscending implements closedts.Storage.
    28  func (s *TestStorage) VisitAscending(nodeID roachpb.NodeID, f func(ctpb.Entry) (done bool)) {
    29  	s.mu.Lock()
    30  	defer s.mu.Unlock()
    31  
    32  	for _, entry := range s.m[nodeID] {
    33  		if f(entry) {
    34  			break
    35  		}
    36  	}
    37  }
    38  
    39  // VisitDescending implements closedts.Storage.
    40  func (s *TestStorage) VisitDescending(nodeID roachpb.NodeID, f func(entry ctpb.Entry) (done bool)) {
    41  	s.mu.Lock()
    42  	defer s.mu.Unlock()
    43  
    44  	for i := len(s.m[nodeID]) - 1; i >= 0; i-- {
    45  		if f(s.m[nodeID][i]) {
    46  			break
    47  		}
    48  	}
    49  }
    50  
    51  // Add implements closedts.Storage.
    52  func (s *TestStorage) Add(nodeID roachpb.NodeID, entry ctpb.Entry) {
    53  	s.mu.Lock()
    54  	defer s.mu.Unlock()
    55  
    56  	if s.m == nil {
    57  		s.m = map[roachpb.NodeID][]ctpb.Entry{}
    58  	}
    59  
    60  	s.m[nodeID] = append(s.m[nodeID], entry)
    61  	sort.Slice(s.m[nodeID], func(i, j int) bool {
    62  		e1, e2 := s.m[nodeID][i], s.m[nodeID][j]
    63  		if e1.ClosedTimestamp == e2.ClosedTimestamp {
    64  			return e1.Epoch < e2.Epoch
    65  		}
    66  		return e1.ClosedTimestamp.Less(e2.ClosedTimestamp)
    67  	})
    68  }
    69  
    70  // Clear implements closedts.Storage.
    71  func (s *TestStorage) Clear() {
    72  	s.mu.Lock()
    73  	defer s.mu.Unlock()
    74  	s.m = nil
    75  }
    76  
    77  // Snapshot returns a copy of the data contain within the TestStorage.
    78  func (s *TestStorage) Snapshot() map[roachpb.NodeID][]ctpb.Entry {
    79  	s.mu.Lock()
    80  	defer s.mu.Unlock()
    81  
    82  	m := map[roachpb.NodeID][]ctpb.Entry{}
    83  	for nodeID, entries := range s.m {
    84  		m[nodeID] = append([]ctpb.Entry(nil), entries...)
    85  	}
    86  	return m
    87  }