improved agent tracking (#21)

go.mod

@@ -32,7 +32,10 @@ require (
 
 )
 
-require github.com/dennis-tra/nebula-crawler v0.0.0-20241010113859-38e4489a8fa7
+require (
+	github.com/dennis-tra/nebula-crawler v0.0.0-20241010113859-38e4489a8fa7
+	github.com/patrickmn/go-cache v2.1.0+incompatible
+)
 
 require (
 	github.com/btcsuite/btcd/btcec/v2 v2.3.4 // indirect
@@ -153,7 +156,7 @@ require (
 	github.com/pkg/errors v0.9.1 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/polydawn/refmt v0.89.0 // indirect
-	github.com/prometheus/client_golang v1.20.4 // indirect
+	github.com/prometheus/client_golang v1.20.4
 	github.com/prometheus/client_model v0.6.1 // indirect
 	github.com/prometheus/common v0.60.0 // indirect
 	github.com/prometheus/procfs v0.15.1 // indirect

go.sum

@@ -632,6 +632,8 @@ github.com/oschwald/maxminddb-golang v1.13.1 h1:G3wwjdN9JmIK2o/ermkHM+98oX5fS+k5
 github.com/oschwald/maxminddb-golang v1.13.1/go.mod h1:K4pgV9N/GcK694KSTmVSDTODk4IsCNThNdTmnaBZ/F8=
 github.com/pascaldekloe/goe v0.0.0-20180627143212-57f6aae5913c/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
 github.com/pascaldekloe/goe v0.1.0/go.mod h1:lzWF7FIEvWOWxwDKqyGYQf6ZUaNfKdP144TG7ZOy1lc=
+github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
+github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 h1:onHthvaw9LFnH4t2DcNVpwGmV9E1BkGknEliJkfwQj0=
 github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58/go.mod h1:DXv8WO4yhMYhSNPKjeNKa5WY9YCIEBRbNzFFPJbWO6Y=
 github.com/pelletier/go-toml v1.9.5/go.mod h1:u1nR/EPcESfeI/szUZKdtJ0xRNbUoANCkoOuaOx1Y+c=

queen.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/dennis-tra/nebula-crawler/maxmind"
 	"github.com/dennis-tra/nebula-crawler/udger"
+	"github.com/patrickmn/go-cache"
 	"github.com/probe-lab/ants-watch/db"
 	"github.com/probe-lab/ants-watch/db/models"
 	tele "github.com/probe-lab/ants-watch/metrics"
@@ -38,8 +39,9 @@ type Queen struct {
 	nebulaDB *NebulaDB
 	keysDB   *KeysDB
 
-	peerstore peerstore.Peerstore
-	datastore ds.Batching
+	peerstore   peerstore.Peerstore
+	datastore   ds.Batching
+	agentsCache *cache.Cache
 
 	ants     []*Ant
 	antsLogs chan antslog.RequestLog
@@ -79,6 +81,7 @@ func NewQueen(ctx context.Context, dbConnString string, keysDbPath string, nPort
 		datastore:        dssync.MutexWrap(ds.NewMapDatastore()),
 		ants:             []*Ant{},
 		antsLogs:         make(chan antslog.RequestLog, 1024),
+		agentsCache:      cache.New(4*24*time.Hour, time.Hour), // 4 days of cache, clean every hour
 		upnp:             true,
 		dbc:              getDbClient(ctx),
 		mmc:              mmc,
@@ -248,9 +251,14 @@ func (q *Queen) consumeAntsLogs(ctx context.Context) {
 			var agent string
 			peerstoreAgent, err := q.peerstore.Get(log.Requester, "AgentVersion")
 			if err != nil {
-				agent = ""
+				if peerstoreAgent, ok := q.agentsCache.Get(log.Requester.String()); ok {
+					agent = peerstoreAgent.(string)
+				} else {
+					agent = ""
+				}
 			} else {
 				agent = peerstoreAgent.(string)
+				q.agentsCache.Set(log.Requester.String(), agent, 0)
 			}
 
 			protocols, _ := q.peerstore.GetProtocols(log.Requester)