repo_description: An open-source authorization as a service inspired by Google Zanzibar, designed to build and manage fine-grained and scalable authorization systems for any application. — Permify is now part of FusionAuth 🎉
repo_topics: abac access-control acl authorization ciam cloud-native distributed fga fine-grained-access-control golang grpc kubernetes least-privilege permission permissions rbac rebac security zanzibar
code: package balancer
  
  import (
  	"errors"
  	"fmt"
  	"log/slog"
  
  	"google.golang.org/grpc/balancer"
  	"google.golang.org/grpc/balancer/base"
  	"google.golang.org/grpc/connectivity"
  	"google.golang.org/grpc/resolver"
  
  	"github.com/Permify/permify/pkg/consistent"
  )
  
  type Balancer struct {
  	// Current overall connectivity state of the balancer.
  	state connectivity.State
  
  	// The ClientConn to communicate with the gRPC client.
  	clientConn ClientConnWrapper
  
  	// Current picker used to select SubConns for requests.
  	picker balancer.Picker
  
  	// Evaluates connectivity state transitions for SubConns.
  	connectivityEvaluator *balancer.ConnectivityStateEvaluator
  
  	// Map of resolver addresses to SubConns.
  	addressSubConns *resolver.AddressMap
  
  	// Tracks the connectivity state of each SubConn.
  	subConnStates map[balancer.SubConn]connectivity.State
  
  	// Configuration for consistent hashing and replication.
  	config *Config
  
  	// Consistent hashing mechanism to distribute requests.
  	consistent *consistent.Consistent
  
  	// Hasher used by the consistent hashing mechanism.
  	hasher consistent.Hasher
  
  	// Stores the last resolver error encountered.
  	lastResolverError error
  
  	// Stores the last connection error encountered.
  	lastConnectionError error
  }
  
  func (b *Balancer) ResolverError(err error) {
  	b.lastResolverError = err
  	if b.addressSubConns.Len() == 0 {
  		b.state = connectivity.TransientFailure
  		b.picker = base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError))
  	}
  
  	if b.state != connectivity.TransientFailure {
  		return
  	}
  
  	// Update the balancer state and picker.
  	b.clientConn.UpdateState(balancer.State{
  		ConnectivityState: b.state,
  		Picker:            b.picker,
  	})
  }
  
  func (b *Balancer) UpdateClientConnState(s balancer.ClientConnState) error {
  	// Log the new ClientConn state.
  	slog.Info("Received new ClientConn state",
  		slog.Any("state", s),
  	)
  
  	// Reset any existing resolver error.
  	b.lastResolverError = nil
  
  	// Handle changes to the balancer configuration.
  	if s.BalancerConfig != nil {
  		svcConfig := s.BalancerConfig.(*Config)
  		if b.config == nil || svcConfig.ReplicationFactor != b.config.ReplicationFactor {
  			slog.Info("Updating consistent hashing configuration",
  				slog.Int("partition_count", svcConfig.PartitionCount),
  				slog.Int("replication_factor", svcConfig.ReplicationFactor),
  				slog.Float64("load", svcConfig.Load),
  				slog.Int("picker_width", svcConfig.PickerWidth),
  			)
  			b.consistent = consistent.New(consistent.Config{
  				PartitionCount:    svcConfig.PartitionCount,
  				ReplicationFactor: svcConfig.ReplicationFactor,
  				Load:              svcConfig.Load,
  				PickerWidth:       svcConfig.PickerWidth,
  				Hasher:            b.hasher,
  			})
  			b.config = svcConfig
  		}
  	}
  
  	// Check if the consistent hashing configuration exists.
  	if b.consistent == nil {
  		slog.Error("No consistent hashing configuration found")
  		b.picker = base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError))
  		b.clientConn.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
  		return fmt.Errorf("no consistent hashing configuration found")
  	}
  
  	// Maintain a set of addresses provided by the resolver.
  	addrsSet := resolver.NewAddressMap()
  	for _, addr := range s.ResolverState.Addresses {
  		addrsSet.Set(addr, nil)
  
  		// Add new SubConns for addresses that are not already tracked.
  		if _, ok := b.addressSubConns.Get(addr); !ok {
  			sc, err := b.clientConn.NewSubConn([]resolver.Address{addr}, balancer.NewSubConnOptions{HealthCheckEnabled: false})
  			if err != nil {
  				slog.Warn("Failed to create new SubConn",
  					slog.String("address", addr.Addr),
  					slog.String("server_name", addr.ServerName),
  					slog.String("error", err.Error()),
  				)
  				continue
  			}
  
  			b.addressSubConns.Set(addr, sc)
  			b.subConnStates[sc] = connectivity.Idle
  			b.connectivityEvaluator.RecordTransition(connectivity.Shutdown, connectivity.Idle)
  			sc.Connect()
  
  			b.consistent.Add(ConsistentMember{
  				SubConn: sc,
  				name:    fmt.Sprintf("%s|%s", addr.ServerName, addr.Addr),
  			})
  		}
  	}
  
  	// Remove SubConns that are no longer part of the resolved addresses.
  	for _, addr := range b.addressSubConns.Keys() {
  		sci, _ := b.addressSubConns.Get(addr)
  		sc := sci.(balancer.SubConn)
  		if _, ok := addrsSet.Get(addr); !ok {
  			slog.Info("Removing SubConn",
  				slog.String("address", addr.Addr),
  				slog.String("server_name", addr.ServerName),
  			)
  			b.clientConn.RemoveSubConn(sc)
  			b.addressSubConns.Delete(addr)
  			b.consistent.Remove(ConsistentMember{
  				SubConn: sc,
  				name:    fmt.Sprintf("%s|%s", addr.ServerName, addr.Addr),
  			}.String())
  		}
  	}
  
  	// Log the current members in the consistent hashing ring.
  	slog.Info("Current consistent members",
  		slog.Int("member_count", len(b.consistent.Members())),
  	)
  	for _, m := range b.consistent.Members() {
  		slog.Info("Consistent member", slog.String("member", m.String()))
  	}
  
  	// Handle the case where the resolver produces zero addresses.
  	if len(s.ResolverState.Addresses) == 0 {
  		err := errors.New("resolver produced zero addresses")
  		b.ResolverError(err)
  		slog.Error("Resolver produced zero addresses")
  		return balancer.ErrBadResolverState
  	}
  
  	// Update the picker based on the current balancer state.
  	if b.state == connectivity.TransientFailure {
  		slog.Warn("Transient failure detected, using error picker")
  		b.picker = base.NewErrPicker(errors.Join(b.lastConnectionError, b.lastResolverError))
  	} else {
  		width := b.config.PickerWidth
  		if width < 1 {
  			width = 1
  		}
  		slog.Info("Creating new picker",
  			slog.Int("width", width),
  		)
  		b.picker = &picker{
  			consistent: b.consistent,
  			width:      width,
  		}
  	}
  
  	// Update the ClientConn state with the new picker.
  	slog.Info("Updating ClientConn state",
  		slog.String("connectivity_state", b.state.String()),
  	)
  	b.clientConn.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
  
  	return nil
  }
  
  func (b *Balancer) UpdateSubConnState(sc balancer.SubConn, state balancer.SubConnState) {
  	s := state.ConnectivityState
  	slog.Info("Received SubConn state change",
  		slog.String("connectivity_state", s.String()),
  		slog.String("sub_conn", fmt.Sprintf("%p", sc)),
  	)
  
  	oldS, ok := b.subConnStates[sc]
  	if !ok {
  		slog.Warn("State change for unknown SubConn",
  			slog.String("connectivity_state", s.String()),
  			slog.String("sub_conn", fmt.Sprintf("%p", sc)),
  		)
  		return
  	}
  
  	if oldS == connectivity.TransientFailure && (s == connectivity.Connecting || s == connectivity.Idle) {
  		if s == connectivity.Idle {
  			slog.Info("Transitioning SubConn to connecting state",
  				slog.String("sub_conn", fmt.Sprintf("%p", sc)),
  			)
  			sc.Connect()
  		}
  		return
  	}
  
  	b.subConnStates[sc] = s
  	switch s {
  	case connectivity.Idle:
  		slog.Info("SubConn is idle, initiating connection",
  			slog.String("sub_conn", fmt.Sprintf("%p", sc)),
  		)
  		sc.Connect()
  	case connectivity.Shutdown:
  		slog.Info("Removing shutdown SubConn",
  			slog.String("sub_conn", fmt.Sprintf("%p", sc)),
  		)
  		delete(b.subConnStates, sc)
  	case connectivity.TransientFailure:
  		slog.Warn("SubConn in transient failure",
  			slog.String("sub_conn", fmt.Sprintf("%p", sc)),
  			slog.String("error", state.ConnectionError.Error()),
  		)
  		b.lastConnectionError = state.ConnectionError
  	}
  
  	b.state = b.connectivityEvaluator.RecordTransition(oldS, s)
  	slog.Info("Updating ClientConn state",
  		slog.String("connectivity_state", b.state.String()),
  	)
  	b.clientConn.UpdateState(balancer.State{ConnectivityState: b.state, Picker: b.picker})
  }
  
  func (b *Balancer) Close() {}
  
  // ExitIdle instructs the LB policy to reconnect to backends / exit the
  // IDLE state, if appropriate and possible. Note that SubConns that enter
  // the IDLE state will not reconnect until SubConn.Connect is called.
  func (b *Balancer) ExitIdle() {
  	// For this consistent hashing balancer, we can attempt to reconnect
  	// idle SubConns by calling Connect() on them
  	for sc, state := range b.subConnStates {
  		if state == connectivity.Idle {
  			sc.Connect()
  		}
  	}
  }
  