open-consul/agent/consul/server_lookup.go

80 lines
2.0 KiB
Go

// Copyright (c) HashiCorp, Inc.
// SPDX-License-Identifier: MPL-2.0
package consul
import (
"fmt"
"sync"
"github.com/hashicorp/consul/agent/metadata"
"github.com/hashicorp/raft"
)
// ServerLookup encapsulates looking up servers by id and address
type ServerLookup struct {
lock sync.RWMutex
addressToServer map[raft.ServerAddress]*metadata.Server
idToServer map[raft.ServerID]*metadata.Server
}
func NewServerLookup() *ServerLookup {
return &ServerLookup{
addressToServer: make(map[raft.ServerAddress]*metadata.Server),
idToServer: make(map[raft.ServerID]*metadata.Server),
}
}
func (sl *ServerLookup) AddServer(server *metadata.Server) {
sl.lock.Lock()
defer sl.lock.Unlock()
sl.addressToServer[raft.ServerAddress(server.Addr.String())] = server
sl.idToServer[raft.ServerID(server.ID)] = server
}
func (sl *ServerLookup) RemoveServer(server *metadata.Server) {
sl.lock.Lock()
defer sl.lock.Unlock()
delete(sl.addressToServer, raft.ServerAddress(server.Addr.String()))
delete(sl.idToServer, raft.ServerID(server.ID))
}
// Implements the ServerAddressProvider interface
func (sl *ServerLookup) ServerAddr(id raft.ServerID) (raft.ServerAddress, error) {
sl.lock.RLock()
defer sl.lock.RUnlock()
svr, ok := sl.idToServer[id]
if !ok {
return "", fmt.Errorf("Could not find address for server id %v", id)
}
return raft.ServerAddress(svr.Addr.String()), nil
}
// Server looks up the server by address, returns a boolean if not found
func (sl *ServerLookup) Server(addr raft.ServerAddress) *metadata.Server {
sl.lock.RLock()
defer sl.lock.RUnlock()
return sl.addressToServer[addr]
}
func (sl *ServerLookup) Servers() []*metadata.Server {
sl.lock.RLock()
defer sl.lock.RUnlock()
var ret []*metadata.Server
for _, svr := range sl.addressToServer {
ret = append(ret, svr)
}
return ret
}
func (sl *ServerLookup) CheckServers(fn func(srv *metadata.Server) bool) {
sl.lock.RLock()
defer sl.lock.RUnlock()
for _, srv := range sl.addressToServer {
if !fn(srv) {
return
}
}
}