feat: add ability to exclude a list of nodes when querying

This commit is contained in:
Derrick Hammer 2024-02-27 07:27:10 -05:00
parent 82de843ad9
commit e9f4a7b0b9
Signed by: pcfreak30
GPG Key ID: C997C339BE476FF2
1 changed files with 10 additions and 0 deletions

View File

@ -26,6 +26,7 @@ type StorageLocationProviderImpl struct {
isWaitingForUri bool isWaitingForUri bool
mutex sync.Mutex mutex sync.Mutex
logger *zap.Logger logger *zap.Logger
excludeNodes []*encoding.NodeId
} }
func (s *StorageLocationProviderImpl) Start() error { func (s *StorageLocationProviderImpl) Start() error {
@ -43,6 +44,10 @@ func (s *StorageLocationProviderImpl) Start() error {
continue continue
} }
if containsNode(s.excludeNodes, nodeId) {
continue
}
s.availableNodes = append(s.availableNodes, nodeId) s.availableNodes = append(s.availableNodes, nodeId)
} }
@ -91,6 +96,9 @@ func (s *StorageLocationProviderImpl) Start() error {
s.logger.Error("Error decoding node id", zap.Error(err)) s.logger.Error("Error decoding node id", zap.Error(err))
continue continue
} }
if containsNode(s.excludeNodes, nodeId) {
continue
}
if !containsNode(s.availableNodes, nodeId) { if !containsNode(s.availableNodes, nodeId) {
s.availableNodes = append(s.availableNodes, nodeId) s.availableNodes = append(s.availableNodes, nodeId)
hasNewNode = true hasNewNode = true
@ -180,6 +188,7 @@ func NewStorageLocationProvider(params StorageLocationProviderParams) *StorageLo
timeoutDuration: 60 * time.Second, timeoutDuration: 60 * time.Second,
uris: make(map[string]storage.StorageLocation), uris: make(map[string]storage.StorageLocation),
logger: params.Logger, logger: params.Logger,
excludeNodes: params.ExcludeNodes,
} }
} }
func containsNode(slice []*encoding.NodeId, item *encoding.NodeId) bool { func containsNode(slice []*encoding.NodeId, item *encoding.NodeId) bool {
@ -195,5 +204,6 @@ type StorageLocationProviderParams struct {
Services service.Services Services service.Services
Hash *encoding.Multihash Hash *encoding.Multihash
LocationTypes []types.StorageLocationType LocationTypes []types.StorageLocationType
ExcludeNodes []*encoding.NodeId
service.ServiceParams service.ServiceParams
} }