trpc.group/trpc-go/trpc-go@v1.0.3/naming/discovery/discovery.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  // Package discovery is a pluggable service discovery module.
    15  package discovery
    16  
    17  import (
    18  	"sync"
    19  
    20  	"trpc.group/trpc-go/trpc-go/naming/registry"
    21  )
    22  
    23  // DefaultDiscovery is the default discovery determined by configuration file.
    24  var DefaultDiscovery Discovery = &IPDiscovery{}
    25  
    26  // SetDefaultDiscovery sets the default discovery.
    27  func SetDefaultDiscovery(d Discovery) {
    28  	DefaultDiscovery = d
    29  }
    30  
    31  // Discovery is the interface that returns nodes by service name.
    32  type Discovery interface {
    33  	List(serviceName string, opt ...Option) (nodes []*registry.Node, err error)
    34  }
    35  
    36  var (
    37  	discoveries = make(map[string]Discovery)
    38  	lock        = sync.RWMutex{}
    39  )
    40  
    41  // Register registers a named discovery.
    42  func Register(name string, s Discovery) {
    43  	lock.Lock()
    44  	discoveries[name] = s
    45  	lock.Unlock()
    46  }
    47  
    48  // Get gets a named discovery.
    49  func Get(name string) Discovery {
    50  	lock.RLock()
    51  	d := discoveries[name]
    52  	lock.RUnlock()
    53  	return d
    54  }
    55  
    56  func unregisterForTesting(name string) {
    57  	lock.Lock()
    58  	delete(discoveries, name)
    59  	lock.Unlock()
    60  }