From f23e4279c42aec766eb6a89562c1fed3a1b97e09 Mon Sep 17 00:00:00 2001 From: Simon Law Date: Sun, 13 Jul 2025 05:47:56 -0700 Subject: [PATCH] types/lazy: add lazy.GMap: a map of lazily computed GValues (#16532) Fixes tailscale/corp#30360 Signed-off-by: Simon Law --- cmd/stund/depaware.txt | 2 +- types/lazy/map.go | 62 +++++++++++++++++++++++++++ types/lazy/map_test.go | 95 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 types/lazy/map.go create mode 100644 types/lazy/map_test.go diff --git a/cmd/stund/depaware.txt b/cmd/stund/depaware.txt index da7680394..81544b750 100644 --- a/cmd/stund/depaware.txt +++ b/cmd/stund/depaware.txt @@ -76,7 +76,7 @@ tailscale.com/cmd/stund dependencies: (generated by github.com/tailscale/depawar L 💣 tailscale.com/util/dirwalk from tailscale.com/metrics tailscale.com/util/dnsname from tailscale.com/tailcfg tailscale.com/util/lineiter from tailscale.com/version/distro - tailscale.com/util/mak from tailscale.com/syncs + tailscale.com/util/mak from tailscale.com/syncs+ tailscale.com/util/nocasemaps from tailscale.com/types/ipproto tailscale.com/util/rands from tailscale.com/tsweb tailscale.com/util/slicesx from tailscale.com/tailcfg diff --git a/types/lazy/map.go b/types/lazy/map.go new file mode 100644 index 000000000..75a1dd739 --- /dev/null +++ b/types/lazy/map.go @@ -0,0 +1,62 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +package lazy + +import "tailscale.com/util/mak" + +// GMap is a map of lazily computed [GValue] pointers, keyed by a comparable +// type. +// +// Use either Get or GetErr, depending on whether your fill function returns an +// error. +// +// GMap is not safe for concurrent use. +type GMap[K comparable, V any] struct { + store map[K]*GValue[V] +} + +// Len returns the number of entries in the map. +func (s *GMap[K, V]) Len() int { + return len(s.store) +} + +// Set attempts to set the value of k to v, and reports whether it succeeded. +// Set only succeeds if k has never been called with Get/GetErr/Set before. +func (s *GMap[K, V]) Set(k K, v V) bool { + z, ok := s.store[k] + if !ok { + z = new(GValue[V]) + mak.Set(&s.store, k, z) + } + return z.Set(v) +} + +// MustSet sets the value of k to v, or panics if k already has a value. +func (s *GMap[K, V]) MustSet(k K, v V) { + if !s.Set(k, v) { + panic("Set after already filled") + } +} + +// Get returns the value for k, computing it with fill if it's not already +// present. +func (s *GMap[K, V]) Get(k K, fill func() V) V { + z, ok := s.store[k] + if !ok { + z = new(GValue[V]) + mak.Set(&s.store, k, z) + } + return z.Get(fill) +} + +// GetErr returns the value for k, computing it with fill if it's not already +// present. +func (s *GMap[K, V]) GetErr(k K, fill func() (V, error)) (V, error) { + z, ok := s.store[k] + if !ok { + z = new(GValue[V]) + mak.Set(&s.store, k, z) + } + return z.GetErr(fill) +} diff --git a/types/lazy/map_test.go b/types/lazy/map_test.go new file mode 100644 index 000000000..ec1152b0b --- /dev/null +++ b/types/lazy/map_test.go @@ -0,0 +1,95 @@ +// Copyright (c) Tailscale Inc & AUTHORS +// SPDX-License-Identifier: BSD-3-Clause + +package lazy + +import ( + "errors" + "testing" +) + +func TestGMap(t *testing.T) { + var gm GMap[string, int] + n := int(testing.AllocsPerRun(1000, func() { + got := gm.Get("42", fortyTwo) + if got != 42 { + t.Fatalf("got %v; want 42", got) + } + })) + if n != 0 { + t.Errorf("allocs = %v; want 0", n) + } +} + +func TestGMapErr(t *testing.T) { + var gm GMap[string, int] + n := int(testing.AllocsPerRun(1000, func() { + got, err := gm.GetErr("42", func() (int, error) { + return 42, nil + }) + if got != 42 || err != nil { + t.Fatalf("got %v, %v; want 42, nil", got, err) + } + })) + if n != 0 { + t.Errorf("allocs = %v; want 0", n) + } + + var gmErr GMap[string, int] + wantErr := errors.New("test error") + n = int(testing.AllocsPerRun(1000, func() { + got, err := gmErr.GetErr("42", func() (int, error) { + return 0, wantErr + }) + if got != 0 || err != wantErr { + t.Fatalf("got %v, %v; want 0, %v", got, err, wantErr) + } + })) + if n != 0 { + t.Errorf("allocs = %v; want 0", n) + } +} + +func TestGMapSet(t *testing.T) { + var gm GMap[string, int] + if !gm.Set("42", 42) { + t.Fatalf("Set failed") + } + if gm.Set("42", 43) { + t.Fatalf("Set succeeded after first Set") + } + n := int(testing.AllocsPerRun(1000, func() { + got := gm.Get("42", fortyTwo) + if got != 42 { + t.Fatalf("got %v; want 42", got) + } + })) + if n != 0 { + t.Errorf("allocs = %v; want 0", n) + } +} + +func TestGMapMustSet(t *testing.T) { + var gm GMap[string, int] + gm.MustSet("42", 42) + defer func() { + if e := recover(); e == nil { + t.Errorf("unexpected success; want panic") + } + }() + gm.MustSet("42", 43) +} + +func TestGMapRecursivePanic(t *testing.T) { + defer func() { + if e := recover(); e != nil { + t.Logf("got panic, as expected") + } else { + t.Errorf("unexpected success; want panic") + } + }() + gm := GMap[string, int]{} + gm.Get("42", func() int { + return gm.Get("42", func() int { return 42 }) + }) +}