micro/cluster/hasql/picker.go

package sql

import (
	"fmt"
	"math"
	"time"

	"golang.yandex/hasql/v2"
)

// compile time guard
var _ hasql.NodePicker[Querier] = (*CustomPicker[Querier])(nil)

// CustomPickerOptions holds options to pick nodes
type CustomPickerOptions struct {
	MaxLag   int
	Priority map[string]int32
	Retries  int
}

// CustomPickerOption func apply option to CustomPickerOptions
type CustomPickerOption func(*CustomPickerOptions)

// CustomPickerMaxLag specifies max lag for which node can be used
func CustomPickerMaxLag(n int) CustomPickerOption {
	return func(o *CustomPickerOptions) {
		o.MaxLag = n
	}
}

// NewCustomPicker creates new node picker
func NewCustomPicker[T Querier](opts ...CustomPickerOption) *CustomPicker[Querier] {
	options := CustomPickerOptions{}
	for _, o := range opts {
		o(&options)
	}
	return &CustomPicker[Querier]{opts: options}
}

// CustomPicker holds node picker options
type CustomPicker[T Querier] struct {
	opts CustomPickerOptions
}

// PickNode used to return specific node
func (p *CustomPicker[T]) PickNode(cnodes []hasql.CheckedNode[T]) hasql.CheckedNode[T] {
	for _, n := range cnodes {
		fmt.Printf("node %s\n", n.Node.String())
	}
	return cnodes[0]
}

func (p *CustomPicker[T]) getPriority(nodeName string) int32 {
	if prio, ok := p.opts.Priority[nodeName]; ok {
		return prio
	}
	return math.MaxInt32 // Default to lowest priority
}

// CompareNodes used to sort nodes
func (p *CustomPicker[T]) CompareNodes(a, b hasql.CheckedNode[T]) int {
	// Get replication lag values
	aLag := a.Info.(interface{ ReplicationLag() int }).ReplicationLag()
	bLag := b.Info.(interface{ ReplicationLag() int }).ReplicationLag()

	// First check that lag lower then MaxLag
	if aLag > p.opts.MaxLag && bLag > p.opts.MaxLag {
		return 0 // both are equal
	}

	// If one node exceeds MaxLag and the other doesn't, prefer the one that doesn't
	if aLag > p.opts.MaxLag {
		return 1 // b is better
	}
	if bLag > p.opts.MaxLag {
		return -1 // a is better
	}

	// Get node priorities
	aPrio := p.getPriority(a.Node.String())
	bPrio := p.getPriority(b.Node.String())

	// if both priority equals
	if aPrio == bPrio {
		// First compare by replication lag
		if aLag < bLag {
			return -1
		}
		if aLag > bLag {
			return 1
		}
		// If replication lag is equal, compare by latency
		aLatency := a.Info.(interface{ Latency() time.Duration }).Latency()
		bLatency := b.Info.(interface{ Latency() time.Duration }).Latency()

		if aLatency < bLatency {
			return -1
		}
		if aLatency > bLatency {
			return 1
		}

		// If lag and latency is equal
		return 0
	}

	// If priorities are different, prefer the node with lower priority value
	if aPrio < bPrio {
		return -1
	}

	return 1
}