micro/registry/memory/memory.go

163 lines
2.9 KiB
Go
Raw Normal View History

2019-01-14 18:27:25 +03:00
// Package memory provides an in-memory registry
package memory
2015-12-09 22:23:16 +03:00
import (
2019-01-18 20:29:17 +03:00
"context"
2019-01-09 09:24:12 +03:00
"sync"
2019-01-18 20:29:17 +03:00
"time"
2019-01-09 09:24:12 +03:00
2019-01-18 20:29:17 +03:00
"github.com/google/uuid"
2015-12-09 22:23:16 +03:00
"github.com/micro/go-micro/registry"
)
2019-01-14 18:27:25 +03:00
type Registry struct {
2019-01-18 20:29:17 +03:00
options registry.Options
2019-01-09 09:24:12 +03:00
sync.RWMutex
2016-02-26 03:09:06 +03:00
Services map[string][]*registry.Service
2019-01-18 20:29:17 +03:00
Watchers map[string]*Watcher
2016-02-26 03:09:06 +03:00
}
2015-12-09 22:23:16 +03:00
2016-03-14 22:15:42 +03:00
var (
2019-01-18 20:29:17 +03:00
timeout = time.Millisecond * 10
2016-03-14 22:15:42 +03:00
)
/*
2019-01-14 18:27:25 +03:00
// Setup sets mock data
func (m *Registry) Setup() {
2019-01-09 09:24:12 +03:00
m.Lock()
defer m.Unlock()
2019-01-14 18:27:25 +03:00
// add some memory data
m.Services = Data
2016-02-26 03:09:06 +03:00
}
*/
2016-02-26 03:09:06 +03:00
2019-01-18 20:29:17 +03:00
func (m *Registry) watch(r *registry.Result) {
var watchers []*Watcher
m.RLock()
for _, w := range m.Watchers {
watchers = append(watchers, w)
}
m.RUnlock()
for _, w := range watchers {
select {
case <-w.exit:
m.Lock()
delete(m.Watchers, w.id)
m.Unlock()
default:
select {
case w.res <- r:
case <-time.After(timeout):
}
}
}
}
func (m *Registry) Init(opts ...registry.Option) error {
for _, o := range opts {
o(&m.options)
}
// add services
2019-01-09 09:24:12 +03:00
m.Lock()
2019-01-18 20:29:17 +03:00
for k, v := range getServices(m.options.Context) {
s := m.Services[k]
m.Services[k] = addServices(s, v)
}
m.Unlock()
return nil
}
2019-01-09 09:24:12 +03:00
2019-01-18 20:29:17 +03:00
func (m *Registry) Options() registry.Options {
return m.options
}
func (m *Registry) GetService(service string) ([]*registry.Service, error) {
m.RLock()
2016-02-26 03:09:06 +03:00
s, ok := m.Services[service]
if !ok || len(s) == 0 {
2019-01-18 20:29:17 +03:00
m.RUnlock()
2016-02-26 03:09:06 +03:00
return nil, registry.ErrNotFound
}
2019-01-18 20:29:17 +03:00
m.RUnlock()
2016-02-26 03:09:06 +03:00
return s, nil
2015-12-09 22:23:16 +03:00
}
2019-01-14 18:27:25 +03:00
func (m *Registry) ListServices() ([]*registry.Service, error) {
2019-01-18 20:29:17 +03:00
m.RLock()
2016-02-26 03:09:06 +03:00
var services []*registry.Service
for _, service := range m.Services {
services = append(services, service...)
}
2019-01-18 20:29:17 +03:00
m.RUnlock()
2016-02-26 03:09:06 +03:00
return services, nil
2015-12-09 22:23:16 +03:00
}
2019-01-14 18:27:25 +03:00
func (m *Registry) Register(s *registry.Service, opts ...registry.RegisterOption) error {
2019-01-18 20:29:17 +03:00
go m.watch(&registry.Result{Action: "update", Service: s})
2019-01-09 09:24:12 +03:00
2019-01-18 20:29:17 +03:00
m.Lock()
2016-03-14 22:15:42 +03:00
services := addServices(m.Services[s.Name], []*registry.Service{s})
m.Services[s.Name] = services
2019-01-18 20:29:17 +03:00
m.Unlock()
2015-12-09 22:23:16 +03:00
return nil
}
2019-01-14 18:27:25 +03:00
func (m *Registry) Deregister(s *registry.Service) error {
2019-01-18 20:29:17 +03:00
go m.watch(&registry.Result{Action: "delete", Service: s})
2019-01-09 09:24:12 +03:00
2019-01-18 20:29:17 +03:00
m.Lock()
2016-03-14 22:15:42 +03:00
services := delServices(m.Services[s.Name], []*registry.Service{s})
m.Services[s.Name] = services
2019-01-18 20:29:17 +03:00
m.Unlock()
2015-12-09 22:23:16 +03:00
return nil
}
2019-01-14 18:27:25 +03:00
func (m *Registry) Watch(opts ...registry.WatchOption) (registry.Watcher, error) {
2019-01-18 20:29:17 +03:00
var wo registry.WatchOptions
2018-02-19 20:12:37 +03:00
for _, o := range opts {
2019-01-18 20:29:17 +03:00
o(&wo)
}
w := &Watcher{
exit: make(chan bool),
res: make(chan *registry.Result),
id: uuid.New().String(),
wo: wo,
2018-02-19 20:12:37 +03:00
}
2019-01-18 20:29:17 +03:00
m.Lock()
m.Watchers[w.id] = w
m.Unlock()
return w, nil
2015-12-09 22:23:16 +03:00
}
2019-01-14 18:27:25 +03:00
func (m *Registry) String() string {
return "memory"
2015-12-20 00:56:14 +03:00
}
2019-01-18 20:29:17 +03:00
func NewRegistry(opts ...registry.Option) registry.Registry {
options := registry.Options{
Context: context.Background(),
}
2019-01-18 20:29:17 +03:00
for _, o := range opts {
o(&options)
}
services := getServices(options.Context)
if services == nil {
services = make(map[string][]*registry.Service)
}
2017-09-28 13:16:56 +03:00
2019-01-14 18:27:25 +03:00
return &Registry{
2019-01-18 20:29:17 +03:00
options: options,
Services: services,
Watchers: make(map[string]*Watcher),
2019-01-14 18:27:25 +03:00
}
2015-12-09 22:23:16 +03:00
}