Eliminate redundant lock acquisitions in LookupKeyID (#1563)

* Eliminate redundant lock acquisitions in LookupKeyID

Iterate s.keys directly under outer RLock instead of calling
s.Len() and s.Key() which each reacquire the lock.

~2x faster across all set sizes:
  size=1: 13.3→6.0 ns/op
  size=10: 40→21 ns/op (mid), 62→32 ns/op (end)
  size=100: 289→149 ns/op (mid), 566→291 ns/op (end)

* Fix benchmark lint: use fmt.Appendf

Author: lestrrat

jwk/bench_set_test.go

@@ -0,0 +1,54 @@
+package jwk_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/lestrrat-go/jwx/v3/jwk"
+)
+
+func makeKeySet(b *testing.B, n int) jwk.Set {
+	b.Helper()
+	set := jwk.NewSet()
+	for i := range n {
+		key, err := jwk.Import(fmt.Appendf(nil, "secret-key-value-%04d", i))
+		if err != nil {
+			b.Fatalf("failed to create key: %v", err)
+		}
+		if err := key.Set(jwk.KeyIDKey, fmt.Sprintf("kid-%04d", i)); err != nil {
+			b.Fatalf("failed to set key id: %v", err)
+		}
+		if err := set.AddKey(key); err != nil {
+			b.Fatalf("failed to add key: %v", err)
+		}
+	}
+	return set
+}
+
+func BenchmarkLookupKeyID(b *testing.B) {
+	for _, size := range []int{1, 10, 100} {
+		set := makeKeySet(b, size)
+
+		midKID := fmt.Sprintf("kid-%04d", size/2)
+		b.Run(fmt.Sprintf("size=%d/mid", size), func(b *testing.B) {
+			b.ReportAllocs()
+			for range b.N {
+				key, ok := set.LookupKeyID(midKID)
+				if !ok || key == nil {
+					b.Fatal("key not found")
+				}
+			}
+		})
+
+		endKID := fmt.Sprintf("kid-%04d", size-1)
+		b.Run(fmt.Sprintf("size=%d/end", size), func(b *testing.B) {
+			b.ReportAllocs()
+			for range b.N {
+				key, ok := set.LookupKeyID(endKID)
+				if !ok || key == nil {
+					b.Fatal("key not found")
+				}
+			}
+		})
+	}
+}

jwk/set.go

@@ -280,11 +280,7 @@ func (s *set) LookupKeyID(kid string) (Key, bool) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
-	for i := range s.Len() {
-		key, ok := s.Key(i)
-		if !ok {
-			return nil, false
-		}
+	for _, key := range s.keys {
 		gotkid, ok := key.KeyID()
 		if ok && gotkid == kid {
 			return key, true