github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/gossip/node_set.go (about)

     1  // Copyright 2015 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 gossip
    12  
    13  import (
    14  	"context"
    15  
    16  	"github.com/cockroachdb/cockroach/pkg/roachpb"
    17  	"github.com/cockroachdb/cockroach/pkg/util/log"
    18  	"github.com/cockroachdb/cockroach/pkg/util/metric"
    19  )
    20  
    21  // A nodeSet keeps a set of nodes and provides simple node-matched
    22  // management functions. nodeSet is not thread safe.
    23  type nodeSet struct {
    24  	nodes        map[roachpb.NodeID]struct{} // Set of roachpb.NodeID
    25  	placeholders int                         // Number of nodes whose ID we don't know yet.
    26  	maxSize      int                         // Maximum size of set
    27  	gauge        *metric.Gauge               // Gauge for the number of nodes in the set.
    28  }
    29  
    30  func makeNodeSet(maxSize int, gauge *metric.Gauge) nodeSet {
    31  	return nodeSet{
    32  		nodes:   make(map[roachpb.NodeID]struct{}),
    33  		maxSize: maxSize,
    34  		gauge:   gauge,
    35  	}
    36  }
    37  
    38  // hasSpace returns whether there are fewer than maxSize nodes
    39  // in the nodes slice.
    40  func (as nodeSet) hasSpace() bool {
    41  	return as.len() < as.maxSize
    42  }
    43  
    44  // len returns the number of nodes in the set.
    45  func (as nodeSet) len() int {
    46  	return len(as.nodes) + as.placeholders
    47  }
    48  
    49  // asSlice returns the nodes as a slice.
    50  func (as nodeSet) asSlice() []roachpb.NodeID {
    51  	slice := make([]roachpb.NodeID, 0, len(as.nodes))
    52  	for node := range as.nodes {
    53  		slice = append(slice, node)
    54  	}
    55  	return slice
    56  }
    57  
    58  // filter returns a nodeSet containing the nodes which return true when passed
    59  // to the supplied filter function filterFn. filterFn should return true to
    60  // keep a node and false to remove a node. The new nodeSet has a separate
    61  // gauge object from the parent.
    62  func (as nodeSet) filter(filterFn func(node roachpb.NodeID) bool) nodeSet {
    63  	avail := makeNodeSet(as.maxSize,
    64  		metric.NewGauge(metric.Metadata{Name: "TODO(marc)", Help: "TODO(marc)"}))
    65  	for node := range as.nodes {
    66  		if filterFn(node) {
    67  			avail.addNode(node)
    68  		}
    69  	}
    70  	return avail
    71  }
    72  
    73  // hasNode verifies that the supplied node matches a node
    74  // in the slice.
    75  func (as nodeSet) hasNode(node roachpb.NodeID) bool {
    76  	_, ok := as.nodes[node]
    77  	return ok
    78  }
    79  
    80  // setMaxSize adjusts the maximum size allowed for the node set.
    81  func (as *nodeSet) setMaxSize(maxSize int) {
    82  	as.maxSize = maxSize
    83  }
    84  
    85  // addNode adds the node to the nodes set.
    86  func (as *nodeSet) addNode(node roachpb.NodeID) {
    87  	// Account for duplicates by including them in the placeholders tally.
    88  	// We try to avoid duplicate gossip connections, but don't guarantee that
    89  	// they never occur.
    90  	if !as.hasNode(node) {
    91  		as.nodes[node] = struct{}{}
    92  	} else {
    93  		as.placeholders++
    94  	}
    95  	as.updateGauge()
    96  }
    97  
    98  // removeNode removes the node from the nodes set.
    99  func (as *nodeSet) removeNode(node roachpb.NodeID) {
   100  	// Parallel the logic in addNode. If we've already removed the given
   101  	// node ID, it's because we had more than one of them.
   102  	if as.hasNode(node) {
   103  		delete(as.nodes, node)
   104  	} else {
   105  		as.placeholders--
   106  	}
   107  	as.updateGauge()
   108  }
   109  
   110  // addPlaceholder adds another node to the set of tracked nodes, but is
   111  // intended for nodes whose IDs we don't know at the time of adding.
   112  // resolvePlaceholder should be called once we know the ID.
   113  func (as *nodeSet) addPlaceholder() {
   114  	as.placeholders++
   115  	as.updateGauge()
   116  }
   117  
   118  // resolvePlaceholder adds another node to the set of tracked nodes, but is
   119  // intended for nodes whose IDs we don't know at the time of adding.
   120  func (as *nodeSet) resolvePlaceholder(node roachpb.NodeID) {
   121  	as.placeholders--
   122  	as.addNode(node)
   123  }
   124  
   125  func (as *nodeSet) updateGauge() {
   126  	if as.placeholders < 0 {
   127  		log.Fatalf(context.TODO(),
   128  			"nodeSet.placeholders should never be less than 0; gossip logic is broken %+v", as)
   129  	}
   130  	as.gauge.Update(int64(as.len()))
   131  }