feat: implement caching mechanism for IP checks and add runner for concurrent execution

Author: jreisinger

Makefile

@@ -1,5 +1,5 @@
 test:
-	go test ./...
+	go test -race -cover ./...
 
 run: test
 	go run -race ./checkip.go 91.228.166.47

README.md

@@ -99,12 +99,16 @@ You can also use environment variables with the same names.
 
 Data used by some checks is downloaded (cached) to `$HOME/.checkip/` folder. Is gets periodically re-downloaded so it is fresh.
 
+Repeated checks of the same IP address within a single `checkip` process are memoized (cached) in memory, so duplicate inputs in one run reuse already computed results instead of hitting the same providers again.
+
 ## Development
 
 Checkip is easy to extend. If you want to add a new way of checking IP addresses:
 
 1. Write a function of type [check.Func](https://pkg.go.dev/github.com/jreisinger/checkip/check#Func).
-2. Add it to [check.Funcs](https://pkg.go.dev/github.com/jreisinger/checkip/check#Funcs) variable.
+2. Add a [check.Definition](https://pkg.go.dev/github.com/jreisinger/checkip/check#Definition) to [check.Definitions](https://pkg.go.dev/github.com/jreisinger/checkip/check#Definitions).
+
+New checks use process-lifetime memoization by default. If a check must always run live, set its cache policy to `check.CacheNone`.
 
 Typical workflow:
 

check/check.go

@@ -15,27 +15,59 @@ const (
 	InfoAndIsMalicious             // both of the above
 )
 
-// Funcs contains all available functions for checking IP addresses.
-var Funcs = []Func{
-	AbuseIPDB,
-	BlockList,
-	CinsScore,
-	DBip,
-	DnsMX,
-	DnsName,
-	Firehol,
-	IPSum,
-	IPtoASN,
-	IsOnAWS,
-	MaxMind,
-	OTX,
-	Ping,
-	SansISC,
-	Shodan,
-	Spur,
-	Tls,
-	UrlScan,
-	VirusTotal,
+// CachePolicy says whether a check result can be reused within one process.
+type CachePolicy int
+
+const (
+	// CacheProcess reuses the check result for repeated IP addresses within a
+	// single process.
+	CacheProcess CachePolicy = iota
+	// CacheNone always runs the check live.
+	CacheNone
+)
+
+// Definition describes a check and how it should be executed.
+type Definition struct {
+	// Name should be unique across all registered checks.
+	Name string
+	Run  Func
+	// Cache defaults to CacheProcess.
+	Cache CachePolicy
+}
+
+// Definitions contains all available checks and their execution policy.
+var Definitions = []Definition{
+	{Name: "abuseipdb.com", Run: AbuseIPDB},
+	{Name: "blocklist.de", Run: BlockList},
+	{Name: "cinsscore.com", Run: CinsScore},
+	{Name: "db-ip.com", Run: DBip},
+	{Name: "dns MX", Run: DnsMX},
+	{Name: "dns name", Run: DnsName},
+	{Name: "firehol.org", Run: Firehol},
+	{Name: "ipsum.app", Run: IPSum},
+	{Name: "iptoasn.com", Run: IPtoASN},
+	{Name: "is on AWS", Run: IsOnAWS},
+	{Name: "maxmind.com", Run: MaxMind},
+	{Name: "otx.alienvault.com", Run: OTX},
+	{Name: "ping", Run: Ping},
+	{Name: "isc.sans.edu", Run: SansISC},
+	{Name: "shodan.io", Run: Shodan},
+	{Name: "spur.us", Run: Spur},
+	{Name: "tls", Run: Tls},
+	{Name: "urlscan.io", Run: UrlScan},
+	{Name: "virustotal.com", Run: VirusTotal},
+}
+
+// Funcs contains all available check functions, derived from Definitions for
+// backward compatibility.
+var Funcs = funcs(Definitions)
+
+func funcs(definitions []Definition) []Func {
+	funcs := make([]Func, 0, len(definitions))
+	for _, definition := range definitions {
+		funcs = append(funcs, definition.Run)
+	}
+	return funcs
 }
 
 // Type is the type of a Check.

check/check_test.go

@@ -38,6 +38,20 @@ func TestNonEmpty(t *testing.T) {
 	}
 }
 
+func TestDefinitionsHaveUniqueNames(t *testing.T) {
+	seen := make(map[string]struct{}, len(Definitions))
+
+	for _, definition := range Definitions {
+		if definition.Name == "" {
+			t.Fatal("definition name must not be empty")
+		}
+		if _, ok := seen[definition.Name]; ok {
+			t.Fatalf("duplicate definition name %q", definition.Name)
+		}
+		seen[definition.Name] = struct{}{}
+	}
+}
+
 // equal tells whether a and b contain the same elements. A nil argument is
 // equivalent to an empty slice.
 func equal(a, b []string) bool {

checkip.go

@@ -40,6 +40,8 @@ func main() {
 		log.Fatal(err)
 	}
 
+	runner := cli.NewRunner(check.Definitions)
+
 	ipaddrs := make(chan net.IP)
 	results := make(chan Result)
 
@@ -55,7 +57,7 @@ func main() {
 		wg.Add(1)
 		go func() {
 			for ipaddr := range ipaddrs {
-				checks, errors := cli.Run(check.Funcs, ipaddr)
+				checks, errors := runner.Run(ipaddr)
 				for _, e := range errors {
 					log.Print(e)
 				}

cli/cli.go

@@ -10,42 +10,25 @@ import (
 	"net"
 	"os"
 	"sort"
-	"sync"
 
 	"github.com/jreisinger/checkip/check"
 )
 
-// Run runs checks concurrently against the ippaddr.
+// Run runs checks against ipaddr in the supplied order.
 func Run(checkFuncs []check.Func, ipaddr net.IP) (Checks, []error) {
-	var checksMu struct {
-		sync.Mutex
-		checks []check.Check
-	}
-	var errorsMu struct {
-		sync.Mutex
-		errors []error
-	}
+	checks := make(Checks, 0, len(checkFuncs))
+	var errors []error
 
-	var wg sync.WaitGroup
 	for _, cf := range checkFuncs {
-		wg.Add(1)
-		go func(cf check.Func) {
-			defer wg.Done()
-			c, err := cf(ipaddr)
-			if err != nil {
-				errorsMu.Lock()
-				errorsMu.errors = append(errorsMu.errors, err)
-				errorsMu.Unlock()
-				return
-			}
-			checksMu.Lock()
-			checksMu.checks = append(checksMu.checks, c)
-			checksMu.Unlock()
-		}(cf)
+		c, err := cf(ipaddr)
+		if err != nil {
+			errors = append(errors, err)
+			continue
+		}
+		checks = append(checks, c)
 	}
-	wg.Wait()
 
-	return checksMu.checks, errorsMu.errors
+	return checks, errors
 }
 
 type Checks []check.Check

cli/runner.go

@@ -0,0 +1,117 @@
+package cli
+
+import (
+	"net"
+	"sync"
+
+	"github.com/jreisinger/checkip/check"
+)
+
+// Runner executes checks and memoizes cacheable results by IP address. It is
+// safe for concurrent use.
+type Runner struct {
+	cachedDefinitions []check.Definition
+	liveDefinitions   []check.Definition
+
+	mu   sync.Mutex
+	memo map[string]*memoEntry
+}
+
+type memoEntry struct {
+	ready  chan struct{}
+	checks Checks
+	errors []error
+}
+
+// NewRunner creates a runner from check definitions.
+func NewRunner(definitions []check.Definition) *Runner {
+	runner := &Runner{
+		memo: make(map[string]*memoEntry),
+	}
+
+	for _, definition := range definitions {
+		if definition.Cache == check.CacheNone {
+			runner.liveDefinitions = append(runner.liveDefinitions, definition)
+			continue
+		}
+		runner.cachedDefinitions = append(runner.cachedDefinitions, definition)
+	}
+
+	return runner
+}
+
+// Run executes the configured checks for ipaddr.
+func (r *Runner) Run(ipaddr net.IP) (Checks, []error) {
+	checks, errors := r.cachedRun(ipaddr)
+
+	liveChecks, liveErrors := runDefinitions(r.liveDefinitions, ipaddr)
+	checks = append(checks, liveChecks...)
+	errors = append(errors, liveErrors...)
+
+	return checks, errors
+}
+
+func (r *Runner) cachedRun(ipaddr net.IP) (Checks, []error) {
+	if len(r.cachedDefinitions) == 0 {
+		return nil, nil
+	}
+
+	key := ipaddr.String()
+
+	r.mu.Lock()
+	if entry, ok := r.memo[key]; ok {
+		r.mu.Unlock()
+		<-entry.ready
+		return cloneChecks(entry.checks), cloneErrors(entry.errors)
+	}
+
+	entry := &memoEntry{ready: make(chan struct{})}
+	r.memo[key] = entry
+	r.mu.Unlock()
+	defer close(entry.ready)
+
+	checks, errors := runDefinitions(r.cachedDefinitions, ipaddr)
+	entry.checks = cloneChecks(checks)
+	entry.errors = cloneErrors(errors)
+
+	return cloneChecks(checks), cloneErrors(errors)
+}
+
+func runDefinitions(definitions []check.Definition, ipaddr net.IP) (Checks, []error) {
+	checks := make(Checks, 0, len(definitions))
+	var errors []error
+
+	for _, definition := range definitions {
+		result, err := definition.Run(ipaddr)
+		if err != nil {
+			errors = append(errors, err)
+			continue
+		}
+		if result.Description == "" {
+			result.Description = definition.Name
+		}
+		checks = append(checks, result)
+	}
+
+	return checks, errors
+}
+
+func cloneChecks(checks Checks) Checks {
+	if len(checks) == 0 {
+		return nil
+	}
+
+	cloned := make(Checks, len(checks))
+	copy(cloned, checks)
+	return cloned
+}
+
+func cloneErrors(errors []error) []error {
+	if len(errors) == 0 {
+		return nil
+	}
+
+	cloned := make([]error, len(errors))
+	copy(cloned, errors)
+	return cloned
+}