From 4aa96cf0504713e977e1e0dad49ed54ea983fd13 Mon Sep 17 00:00:00 2001
From: Philip Laine <philip.laine@gmail.com>
Date: Wed, 15 May 2024 21:49:37 +0200
Subject: [PATCH] Move mirror metrics code to mirror handler

---
 CHANGELOG.md                 |  1 +
 pkg/registry/distribution.go | 28 ++++++++++++----------
 pkg/registry/registry.go     | 46 ++++++++++++++++++++----------------
 3 files changed, 42 insertions(+), 33 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index bc806dbc..cbd898e4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -20,6 +20,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - [#481](https://github.com/XenitAB/spegel/pull/481) Enable perfsprint linter and fix errors.
 - [#482](https://github.com/XenitAB/spegel/pull/482) Enable gocritic linter and fix errors.
 - [#483](https://github.com/XenitAB/spegel/pull/483) Update errcheck linter configuration and fix errors.
+- [#487](https://github.com/XenitAB/spegel/pull/487) Move mirror metrics code to mirror handler.
 
 ### Deprecated
 
diff --git a/pkg/registry/distribution.go b/pkg/registry/distribution.go
index 8880a335..9a8d0113 100644
--- a/pkg/registry/distribution.go
+++ b/pkg/registry/distribution.go
@@ -17,9 +17,10 @@ const (
 )
 
 type reference struct {
-	kind referenceKind
-	name string
-	dgst digest.Digest
+	kind             referenceKind
+	name             string
+	dgst             digest.Digest
+	originalRegistry string
 }
 
 func (r reference) hasLatestTag() bool {
@@ -43,32 +44,35 @@ var (
 	blobsRegexDigest    = regexp.MustCompile(`/v2/` + nameRegex.String() + `/blobs/(.*)`)
 )
 
-func parsePathComponents(registry, path string) (reference, error) {
+func parsePathComponents(originalRegistry, path string) (reference, error) {
 	comps := manifestRegexTag.FindStringSubmatch(path)
 	if len(comps) == 6 {
-		if registry == "" {
+		if originalRegistry == "" {
 			return reference{}, errors.New("registry parameter needs to be set for tag references")
 		}
-		name := fmt.Sprintf("%s/%s:%s", registry, comps[1], comps[5])
+		name := fmt.Sprintf("%s/%s:%s", originalRegistry, comps[1], comps[5])
 		ref := reference{
-			kind: referenceKindManifest,
-			name: name,
+			kind:             referenceKindManifest,
+			name:             name,
+			originalRegistry: originalRegistry,
 		}
 		return ref, nil
 	}
 	comps = manifestRegexDigest.FindStringSubmatch(path)
 	if len(comps) == 6 {
 		ref := reference{
-			kind: referenceKindManifest,
-			dgst: digest.Digest(comps[5]),
+			kind:             referenceKindManifest,
+			dgst:             digest.Digest(comps[5]),
+			originalRegistry: originalRegistry,
 		}
 		return ref, nil
 	}
 	comps = blobsRegexDigest.FindStringSubmatch(path)
 	if len(comps) == 6 {
 		ref := reference{
-			kind: referenceKindBlob,
-			dgst: digest.Digest(comps[5]),
+			kind:             referenceKindBlob,
+			dgst:             digest.Digest(comps[5]),
+			originalRegistry: originalRegistry,
 		}
 		return ref, nil
 	}
diff --git a/pkg/registry/registry.go b/pkg/registry/registry.go
index 370e46e6..0e1a8713 100644
--- a/pkg/registry/registry.go
+++ b/pkg/registry/registry.go
@@ -170,8 +170,8 @@ func (r *Registry) registryHandler(rw mux.ResponseWriter, req *http.Request) str
 	}
 
 	// Parse out path components from request.
-	registryName := req.URL.Query().Get("ns")
-	ref, err := parsePathComponents(registryName, req.URL.Path)
+	originalRegistry := req.URL.Query().Get("ns")
+	ref, err := parsePathComponents(originalRegistry, req.URL.Path)
 	if err != nil {
 		rw.WriteError(http.StatusNotFound, err)
 		return ""
@@ -182,15 +182,6 @@ func (r *Registry) registryHandler(rw mux.ResponseWriter, req *http.Request) str
 		// Set mirrored header in request to stop infinite loops
 		req.Header.Set(MirroredHeaderKey, "true")
 		r.handleMirror(rw, req, ref)
-		sourceType := "internal"
-		if r.isExternalRequest(req) {
-			sourceType = "external"
-		}
-		cacheType := "hit"
-		if rw.Status() != http.StatusOK {
-			cacheType = "miss"
-		}
-		metrics.MirrorRequestsTotal.WithLabelValues(registryName, cacheType, sourceType).Inc()
 		return "mirror"
 	}
 
@@ -210,12 +201,6 @@ func (r *Registry) registryHandler(rw mux.ResponseWriter, req *http.Request) str
 }
 
 func (r *Registry) handleMirror(rw mux.ResponseWriter, req *http.Request, ref reference) {
-	if !r.resolveLatestTag && ref.hasLatestTag() {
-		r.log.V(4).Info("skipping mirror request for image with latest tag", "image", ref.name)
-		rw.WriteHeader(http.StatusNotFound)
-		return
-	}
-
 	key := ref.dgst.String()
 	if key == "" {
 		key = ref.name
@@ -223,14 +208,33 @@ func (r *Registry) handleMirror(rw mux.ResponseWriter, req *http.Request, ref re
 
 	log := r.log.WithValues("key", key, "path", req.URL.Path, "ip", getClientIP(req))
 
-	// Resolve mirror with the requested key
-	resolveCtx, cancel := context.WithTimeout(req.Context(), r.resolveTimeout)
-	defer cancel()
-	resolveCtx = logr.NewContext(resolveCtx, log)
 	isExternal := r.isExternalRequest(req)
 	if isExternal {
 		log.Info("handling mirror request from external node")
 	}
+
+	defer func() {
+		sourceType := "internal"
+		if isExternal {
+			sourceType = "external"
+		}
+		cacheType := "hit"
+		if rw.Status() != http.StatusOK {
+			cacheType = "miss"
+		}
+		metrics.MirrorRequestsTotal.WithLabelValues(ref.originalRegistry, cacheType, sourceType).Inc()
+	}()
+
+	if !r.resolveLatestTag && ref.hasLatestTag() {
+		r.log.V(4).Info("skipping mirror request for image with latest tag", "image", ref.name)
+		rw.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	// Resolve mirror with the requested key
+	resolveCtx, cancel := context.WithTimeout(req.Context(), r.resolveTimeout)
+	defer cancel()
+	resolveCtx = logr.NewContext(resolveCtx, log)
 	peerCh, err := r.router.Resolve(resolveCtx, key, isExternal, r.resolveRetries)
 	if err != nil {
 		rw.WriteError(http.StatusInternalServerError, err)