Handle discovered ports TTL

Author: silvanocerza

cache.go

@@ -0,0 +1,130 @@
+//
+// This file is part of mdns-discovery.
+//
+// Copyright 2018-2021 ARDUINO SA (http://www.arduino.cc/)
+//
+// This software is released under the GNU General Public License version 3,
+// which covers the main part of arduino-cli.
+// The terms of this license can be found at:
+// https://www.gnu.org/licenses/gpl-3.0.en.html
+//
+// You can be released from the requirements of the above licenses by purchasing
+// a commercial license. Buying such a license is mandatory if you want to modify or
+// otherwise use the software for commercial activities involving the Arduino
+// software without disclosing the source code of your own applications. To purchase
+// a commercial license, send an email to license@arduino.cc.
+//
+
+package main
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	discovery "github.com/arduino/pluggable-discovery-protocol-handler"
+)
+
+// cacheItem stores the Port discovered and its timer to handle TTL.
+type cacheItem struct {
+	Port       *discovery.Port
+	timerTTL   *time.Timer
+	ctx        context.Context
+	cancelFunc context.CancelFunc
+}
+
+// portsCached is a cache to store discovered ports with a TTL.
+// Ports that reach their TTL are automatically deleted and
+// main discovery process is notified.
+// All operations on the internal data are thread safe.
+type portsCache struct {
+	data             map[string]*cacheItem
+	dataMutex        sync.Mutex
+	itemsTTL         time.Duration
+	deletionCallback func(port *discovery.Port)
+}
+
+// NewCache creates a new portsCache and returns it.
+// itemsTTL is the TTL of a single item, when it's reached
+// the stored item is deleted.
+func NewCache(itemsTTL time.Duration) *portsCache {
+	return &portsCache{
+		itemsTTL: itemsTTL,
+		data:     make(map[string]*cacheItem),
+	}
+}
+
+// Set stores a new port and sets its TTL.
+// If the specified key is already found the item's TTL is
+// renewed.
+// Set is thread safe.
+func (c *portsCache) Set(key string, port *discovery.Port) {
+	c.dataMutex.Lock()
+	defer c.dataMutex.Unlock()
+	// We need a cancellable context to avoid leaving
+	// goroutines hanging if an item's timer TTL is stopped.
+	ctx, cancelFunc := context.WithCancel(context.Background())
+	item, ok := c.data[key]
+	timerTTL := time.NewTimer(c.itemsTTL)
+	if ok {
+		item.timerTTL.Stop()
+		// If we stop the timer the goroutine that waits
+		// for it to go off would just hang forever if we
+		// don't cancel the item's context
+		item.cancelFunc()
+		item.timerTTL = timerTTL
+		item.ctx = ctx
+		item.cancelFunc = cancelFunc
+	} else {
+		item = &cacheItem{
+			Port:       port,
+			timerTTL:   timerTTL,
+			ctx:        ctx,
+			cancelFunc: cancelFunc,
+		}
+		c.data[key] = item
+	}
+
+	go func(key string, item *cacheItem) {
+		select {
+		case <-item.timerTTL.C:
+			c.dataMutex.Lock()
+			defer c.dataMutex.Unlock()
+			c.deletionCallback(item.Port)
+			delete(c.data, key)
+			return
+		case <-item.ctx.Done():
+			// The TTL has been renewed, we also stop the timer
+			// that keeps track of the TTL.
+			// The channel of a stopped timer will never fire so
+			// if keep waiting for it in this goroutine it will
+			// hang forever.
+			// Using a context we handle this gracefully.
+			return
+		}
+	}(key, item)
+}
+
+// Get returns the item stored with the specified key and true.
+// If the item is not found returns a nil port and false.
+// Get is thread safe.
+func (c *portsCache) Get(key string) (*discovery.Port, bool) {
+	c.dataMutex.Lock()
+	defer c.dataMutex.Unlock()
+	if item, ok := c.data[key]; ok {
+		return item.Port, ok
+	}
+	return nil, false
+}
+
+// Clear removes all the stored items and stops their TTL timers.
+// Clear is thread safe.
+func (c *portsCache) Clear() {
+	c.dataMutex.Lock()
+	defer c.dataMutex.Unlock()
+	for key, item := range c.data {
+		item.timerTTL.Stop()
+		item.cancelFunc()
+		delete(c.data, key)
+	}
+}

main.go

@@ -40,21 +40,30 @@ func main() {
 		fmt.Fprintf(os.Stderr, "Error: %s\n", err.Error())
 		os.Exit(1)
 	}
-
 }
 
 const mdnsServiceName = "_arduino._tcp"
 
+// Discovered ports stay alive for this amount of time
+// since the last time they've been found by an mDNS query.
+const portsTTL = time.Second * 60
+
+// This is interval at which mDNS queries are made.
+const discoveryInterval = time.Second * 15
+
 // MDNSDiscovery is the implementation of the network pluggable-discovery
 type MDNSDiscovery struct {
 	cancelFunc  func()
 	entriesChan chan *mdns.ServiceEntry
+
+	portsCache *portsCache
 }
 
 // Hello handles the pluggable-discovery HELLO command
 func (d *MDNSDiscovery) Hello(userAgent string, protocolVersion int) error {
 	// The mdns library used has some logs statement that we must disable
 	log.SetOutput(ioutil.Discard)
+	d.portsCache = NewCache(portsTTL)
 	return nil
 }
 
@@ -68,6 +77,9 @@ func (d *MDNSDiscovery) Stop() error {
 		close(d.entriesChan)
 		d.entriesChan = nil
 	}
+	if d.portsCache != nil {
+		d.portsCache.Clear()
+	}
 	return nil
 }
 
@@ -82,20 +94,34 @@ func (d *MDNSDiscovery) StartSync(eventCB discovery.EventCallback, errorCB disco
 		return fmt.Errorf("already syncing")
 	}
 
+	if d.portsCache.deletionCallback == nil {
+		// We can't set the cache deletion callback at creation,
+		// this is the only place we can get the callback
+		d.portsCache.deletionCallback = func(port *discovery.Port) {
+			eventCB("remove", port)
+		}
+	}
+
 	d.entriesChan = make(chan *mdns.ServiceEntry, 4)
 	var receiver <-chan *mdns.ServiceEntry = d.entriesChan
 	var sender chan<- *mdns.ServiceEntry = d.entriesChan
 
 	go func() {
 		for entry := range receiver {
-			eventCB("add", toDiscoveryPort(entry))
+			port := toDiscoveryPort(entry)
+			key := portKey(port)
+			if _, ok := d.portsCache.Get(key); !ok {
+				// Port is not cached so let the user know a new one has been found
+				eventCB("add", port)
+			}
+			d.portsCache.Set(portKey(port), port)
 		}
 	}()
 
 	params := &mdns.QueryParam{
 		Service:             mdnsServiceName,
 		Domain:              "local",
-		Timeout:             time.Second * 15,
+		Timeout:             discoveryInterval,
 		Entries:             sender,
 		WantUnicastResponse: false,
 	}
@@ -117,6 +143,10 @@ func (d *MDNSDiscovery) StartSync(eventCB discovery.EventCallback, errorCB disco
 	return nil
 }
 
+func portKey(p *discovery.Port) string {
+	return fmt.Sprintf("%s:%s %s", p.Address, p.Properties.Get("port"), p.Properties.Get("board"))
+}
+
 func toDiscoveryPort(entry *mdns.ServiceEntry) *discovery.Port {
 	ip := ""
 	if len(entry.AddrV4) > 0 {