Fix host resolution and concurrent cache updates

internal/cache/concurrency_test.go

@@ -0,0 +1,182 @@
+package cache
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestShortcutRejectsAmbiguousTargets(t *testing.T) {
+	c := New("", 3600)
+	c.Upsert(Entry{Name: "web.prod.example", IP: "10.0.0.1"})
+	c.Upsert(Entry{Name: "web.dev.example", IP: "10.0.0.2"})
+	normalize := func(s string) string { return strings.SplitN(s, ".", 2)[0] }
+	for range 100 {
+		entry, found, fresh, err := c.GetByShortcut("web", normalize)
+		if err == nil || found || fresh || entry.Name != "" {
+			t.Fatalf("ambiguous shortcut returned target: %+v, %v, %v, %v", entry, found, fresh, err)
+		}
+	}
+}
+
+func TestSaveMergesIndependentlyLoadedCaches(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "hosts.json")
+	a, b := New(path, 3600), New(path, 3600)
+	for _, c := range []*Cache{a, b} {
+		if err := c.Load(); err != nil {
+			t.Fatal(err)
+		}
+	}
+	a.Upsert(Entry{Name: "a", IP: "10.0.0.1"})
+	b.Upsert(Entry{Name: "b", IP: "10.0.0.2"})
+	for _, c := range []*Cache{a, b} {
+		if err := c.Save(); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := a.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if len(a.All()) != 2 {
+		t.Fatalf("lost an update: %+v", a.All())
+	}
+
+	// Saving a stale usage update must retain newer IPs and both usage times.
+	if err := b.Load(); err != nil {
+		t.Fatal(err)
+	}
+	b.MarkUsed("a")
+	used, _ := b.Get("a")
+	a.Upsert(Entry{Name: "a", IP: "10.0.0.3"})
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.Save(); err != nil {
+		t.Fatal(err)
+	}
+	got, _ := b.Get("a")
+	if got.IP != "10.0.0.3" || !got.LastUsed.Equal(used.LastUsed) {
+		t.Fatalf("lost refreshed IP or usage: %+v", got)
+	}
+
+	b.Delete("a")
+	if err := b.Save(); err != nil {
+		t.Fatal(err)
+	}
+	if err := a.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if got, _ := a.Get("a"); got.Name != "" {
+		t.Fatalf("deleted entry survived: %+v", got)
+	}
+	a.Clear()
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	if err := b.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if len(b.All()) != 0 {
+		t.Fatal("clear merged old entries back in")
+	}
+}
+
+func TestSaveAcrossProcesses(t *testing.T) {
+	if path := os.Getenv("NETSSH_TEST_CACHE_WRITER_PATH"); path != "" {
+		c := New(path, 3600)
+		if err := c.Load(); err != nil {
+			t.Fatal(err)
+		}
+		c.Upsert(Entry{Name: os.Getenv("NETSSH_TEST_CACHE_WRITER_NAME"), IP: "10.0.0.1"})
+		if err := c.Save(); err != nil {
+			t.Fatal(err)
+		}
+		return
+	}
+	path := filepath.Join(t.TempDir(), "hosts.json")
+	const writers = 8
+	cmds := make([]*exec.Cmd, writers)
+	for i := range cmds {
+		cmd := exec.Command(os.Args[0], "-test.run=^TestSaveAcrossProcesses$")
+		cmd.Env = append(os.Environ(), "NETSSH_TEST_CACHE_WRITER_PATH="+path, fmt.Sprintf("NETSSH_TEST_CACHE_WRITER_NAME=host-%d", i))
+		cmd.Stdout, cmd.Stderr = os.Stdout, os.Stderr
+		if err := cmd.Start(); err != nil {
+			t.Fatal(err)
+		}
+		cmds[i] = cmd
+	}
+	for _, cmd := range cmds {
+		if err := cmd.Wait(); err != nil {
+			t.Error(err)
+		}
+	}
+	c := New(path, 3600)
+	if err := c.Load(); err != nil {
+		t.Fatal(err)
+	}
+	if len(c.All()) != writers {
+		t.Fatalf("lost concurrent updates: got %d, want %d", len(c.All()), writers)
+	}
+}
+
+func TestSavePreservesNewerUsageFromDisk(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "hosts.json")
+	a := New(path, 3600)
+	a.Upsert(Entry{Name: "web", IP: "10.0.0.1", LastUsed: time.Now()})
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	b := New(path, 3600)
+	if err := b.Load(); err != nil {
+		t.Fatal(err)
+	}
+	a.MarkUsed("web")
+	used, _ := a.Get("web")
+	if err := a.Save(); err != nil {
+		t.Fatal(err)
+	}
+	b.Upsert(Entry{Name: "web", IP: "10.0.0.2"})
+	if err := b.Save(); err != nil {
+		t.Fatal(err)
+	}
+	got, _ := b.Get("web")
+	if got.IP != "10.0.0.2" || !got.LastUsed.Equal(used.LastUsed) {
+		t.Fatalf("lost data: %+v", got)
+	}
+}
+
+func TestUsageSaveDoesNotRestoreRemovedHosts(t *testing.T) {
+	for _, clear := range []bool{false, true} {
+		t.Run(fmt.Sprintf("clear=%v", clear), func(t *testing.T) {
+			path := filepath.Join(t.TempDir(), "hosts.json")
+			a := New(path, 3600)
+			a.Upsert(Entry{Name: "web", IP: "10.0.0.1"})
+			if err := a.Save(); err != nil {
+				t.Fatal(err)
+			}
+			b := New(path, 3600)
+			if err := b.Load(); err != nil {
+				t.Fatal(err)
+			}
+			b.MarkUsed("web")
+			if clear {
+				a.Clear()
+			} else {
+				a.Delete("web")
+			}
+			if err := a.Save(); err != nil {
+				t.Fatal(err)
+			}
+			if err := b.Save(); err != nil {
+				t.Fatal(err)
+			}
+			if len(b.All()) != 0 {
+				t.Fatalf("usage update restored removed entry: %+v", b.All())
+			}
+		})
+	}
+}