github.com/weaviate/weaviate@v1.24.6/usecases/cluster/iterator.go (about)

     1  //                           _       _
     2  // __      _____  __ ___   ___  __ _| |_ ___
     3  // \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
     4  //  \ V  V /  __/ (_| |\ V /| | (_| | ||  __/
     5  //   \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
     6  //
     7  //  Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
     8  //
     9  //  CONTACT: hello@weaviate.io
    10  //
    11  
    12  package cluster
    13  
    14  import (
    15  	"fmt"
    16  	"math/rand"
    17  )
    18  
    19  type NodeIterationStrategy int
    20  
    21  const (
    22  	StartRandom NodeIterationStrategy = iota
    23  	StartAfter
    24  )
    25  
    26  type NodeIterator struct {
    27  	hostnames []string
    28  	state     int
    29  }
    30  
    31  type HostnameSource interface {
    32  	AllNames() []string
    33  }
    34  
    35  func NewNodeIterator(nodeNames []string,
    36  	strategy NodeIterationStrategy,
    37  ) (*NodeIterator, error) {
    38  	if strategy != StartRandom && strategy != StartAfter {
    39  		return nil, fmt.Errorf("unsupported strategy: %v", strategy)
    40  	}
    41  
    42  	startState := 0
    43  	if strategy == StartRandom {
    44  		startState = rand.Intn(len(nodeNames))
    45  	}
    46  
    47  	return &NodeIterator{
    48  		hostnames: nodeNames,
    49  		state:     startState,
    50  	}, nil
    51  }
    52  
    53  func (n *NodeIterator) SetStartNode(startNode string) {
    54  	for i, node := range n.hostnames {
    55  		if node == startNode {
    56  			n.state = i + 1
    57  			if n.state == len(n.hostnames) {
    58  				n.state = 0
    59  			}
    60  			break
    61  		}
    62  	}
    63  }
    64  
    65  func (n *NodeIterator) Next() string {
    66  	curr := n.hostnames[n.state]
    67  	n.state++
    68  	if n.state == len(n.hostnames) {
    69  		n.state = 0
    70  	}
    71  
    72  	return curr
    73  }