Added Registry TTL to memory registry. Tracking node lifetimes.

registry/handler/handler.go

@@ -2,6 +2,7 @@ package handler
 
 import (
 	"context"
+	"time"
 
 	"github.com/micro/go-micro/errors"
 	"github.com/micro/go-micro/registry"
@@ -26,10 +27,17 @@ func (r *Registry) GetService(ctx context.Context, req *pb.GetRequest, rsp *pb.G
 }
 
 func (r *Registry) Register(ctx context.Context, req *pb.Service, rsp *pb.EmptyResponse) error {
-	err := r.Registry.Register(service.ToService(req))
+	var regOpts []registry.RegisterOption
+	regTTL, ok := ctx.Value("register_ttl").(time.Duration)
+	if ok {
+		regOpts = append(regOpts, registry.RegisterTTL(regTTL))
+	}
+
+	err := r.Registry.Register(service.ToService(req), regOpts...)
 	if err != nil {
 		return errors.InternalServerError("go.micro.registry", err.Error())
 	}
+
 	return nil
 }
 

registry/memory/memory.go

@@ -8,17 +8,25 @@ import (
 
 	"github.com/google/uuid"
 	"github.com/micro/go-micro/registry"
+	"github.com/micro/go-micro/util/log"
 )
 
 var (
 	timeout = time.Millisecond * 10
 )
 
+// node tracks node registration timestamp and TTL
+type node struct {
+	ts  time.Time
+	ttl time.Duration
+}
+
 type Registry struct {
 	options registry.Options
 
 	sync.RWMutex
 	Services map[string][]*registry.Service
+	nodes    map[string]*node
 	Watchers map[string]*Watcher
 }
 
@@ -39,6 +47,7 @@ func NewRegistry(opts ...registry.Option) registry.Registry {
 	return &Registry{
 		options:  options,
 		Services: services,
+		nodes:    make(map[string]*node),
 		Watchers: make(map[string]*Watcher),
 	}
 }
@@ -111,30 +120,73 @@ func (m *Registry) Register(s *registry.Service, opts ...registry.RegisterOption
 	m.Lock()
 	defer m.Unlock()
 
+	var options registry.RegisterOptions
+	for _, o := range opts {
+		o(&options)
+	}
+
 	if service, ok := m.Services[s.Name]; !ok {
 		m.Services[s.Name] = []*registry.Service{s}
+		// add all nodes into nodes map to track their TTL
+		for _, n := range s.Nodes {
+			log.Logf("Tracking node %s for service %s", n.Id, s.Name)
+			m.nodes[s.Name+n.Id] = &node{
+				ts:  time.Now(),
+				ttl: options.TTL,
+			}
+		}
 		go m.sendEvent(&registry.Result{Action: "update", Service: s})
 	} else {
+		// svcCount essentially keep the count of all service vesions
 		svcCount := len(service)
-		svcNodeCounts := make(map[string]map[string]int)
+		// svcNodes maintains a list of node Ids per particular service version
+		svcNodes := make(map[string]map[string][]string)
+		// collect all service ids for all service versions
 		for _, s := range service {
-			if _, ok := svcNodeCounts[s.Name]; !ok {
+			if _, ok := svcNodes[s.Name]; !ok {
-				svcNodeCounts[s.Name] = make(map[string]int)
+				svcNodes[s.Name] = make(map[string][]string)
 			}
-			if _, ok := svcNodeCounts[s.Name][s.Version]; !ok {
+			if _, ok := svcNodes[s.Name][s.Version]; !ok {
-				svcNodeCounts[s.Name][s.Version] = len(s.Nodes)
+				for _, n := range s.Nodes {
+					svcNodes[s.Name][s.Version] = append(svcNodes[s.Name][s.Version], n.Id)
 				}
 			}
-		// if merged count and original service counts changed we added new version of the service
+		}
+		// if merged count and original service counts changed we know we are adding a new version of the service
 		merged := registry.Merge(service, []*registry.Service{s})
 		if len(merged) != svcCount {
 			m.Services[s.Name] = merged
+			// we know s is the new [version of] service; we need to strart tracking its nodes
+			for _, n := range s.Nodes {
+				log.Logf("Tracking node %s for service %s", n.Id, s.Name)
+				m.nodes[s.Name+n.Id] = &node{
+					ts:  time.Now(),
+					ttl: options.TTL,
+				}
+			}
 			go m.sendEvent(&registry.Result{Action: "update", Service: s})
 			return nil
 		}
-		// if the node count for a particular service has changed we added a new node to the service
+		// if the node count of any particular service [version] changed we know we are adding a new node to the service
 		for _, s := range merged {
-			if len(s.Nodes) != svcNodeCounts[s.Name][s.Version] {
+			// we know that if the node counts have changed we need to track new nodes
+			if len(s.Nodes) != len(svcNodes[s.Name][s.Version]) {
+				for _, n := range s.Nodes {
+					var found bool
+					for _, id := range svcNodes[s.Name][s.Version] {
+						if n.Id == id {
+							found = true
+							break
+						}
+					}
+					if !found {
+						log.Logf("Tracking node %s for service %s", n.Id, s.Name)
+						m.nodes[s.Name+n.Id] = &node{
+							ts:  time.Now(),
+							ttl: options.TTL,
+						}
+					}
+				}
 				m.Services[s.Name] = merged
 				go m.sendEvent(&registry.Result{Action: "update", Service: s})
 				return nil
@@ -150,13 +202,45 @@ func (m *Registry) Deregister(s *registry.Service) error {
 	defer m.Unlock()
 
 	if service, ok := m.Services[s.Name]; ok {
-		go m.sendEvent(&registry.Result{Action: "delete", Service: s})
+		// svcNodes maintains a list of node Ids per particular service version
+		svcNodes := make(map[string]map[string][]string)
+		// collect all service ids for all service versions
+		for _, s := range service {
+			if _, ok := svcNodes[s.Name]; !ok {
+				svcNodes[s.Name] = make(map[string][]string)
+			}
+			if _, ok := svcNodes[s.Name][s.Version]; !ok {
+				for _, n := range s.Nodes {
+					svcNodes[s.Name][s.Version] = append(svcNodes[s.Name][s.Version], n.Id)
+				}
+			}
+		}
 		if service := registry.Remove(service, []*registry.Service{s}); len(service) == 0 {
+			id := svcNodes[s.Name][s.Version][0]
+			log.Logf("Stopped tracking node %s for service %s", id, s.Name)
+			delete(m.nodes, s.Name+id)
 			delete(m.Services, s.Name)
 		} else {
+			// find out which nodes have been removed
+			for _, id := range svcNodes[s.Name][s.Version] {
+				for _, s := range service {
+					var found bool
+					for _, n := range s.Nodes {
+						if id == n.Id {
+							found = true
+							break
+						}
+					}
+					if !found {
+						log.Logf("Stopped tracking node %s for service %s", id, s.Name)
+						delete(m.nodes, s.Name+id)
+					}
+				}
 				m.Services[s.Name] = service
 			}
 		}
+		go m.sendEvent(&registry.Result{Action: "delete", Service: s})
+	}
 
 	return nil
 }

registry/service/service.go

@@ -58,8 +58,13 @@ func (s *serviceRegistry) Register(srv *registry.Service, opts ...registry.Regis
 		o(&options)
 	}
 
+	ctx := context.Background()
+	if options.TTL.Nanoseconds() != 0.0 {
+		ctx = context.WithValue(ctx, "register_ttl", options.TTL)
+	}
+
 	// register the service
-	_, err := s.client.Register(context.TODO(), ToProto(srv), s.callOpts()...)
+	_, err := s.client.Register(ctx, ToProto(srv), s.callOpts()...)
 	if err != nil {
 		return err
 	}