Basics of Cert-Count Non-Locking Telemetry

builtin/logical/pki/backend.go

@@ -3,11 +3,14 @@ package pki
 import (
 	"context"
 	"fmt"
+	"sort"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
 
+	atomic2 "go.uber.org/atomic"
+
 	"github.com/hashicorp/vault/sdk/helper/consts"
 
 	"github.com/armon/go-metrics"
@@ -198,6 +201,13 @@ func Backend(conf *logical.BackendConfig) *backend {
 	// Delay the first tidy until after we've started up.
 	b.lastTidy = time.Now()
 
+	// Metrics initialization for count of certificates in storage
+	b.certsCounted = atomic2.NewBool(false)
+	b.certCount = new(uint32)
+	b.revokedCertCount = new(uint32)
+	b.possibleDoubleCountedSerials = make([]string, 0, 250)
+	b.possibleDoubleCountedRevokedSerials = make([]string, 0, 250)
+
 	return &b
 }
 
@@ -214,6 +224,12 @@ type backend struct {
 	tidyStatus     *tidyStatus
 	lastTidy       time.Time
 
+	certCount                           *uint32
+	revokedCertCount                    *uint32
+	certsCounted                        *atomic2.Bool
+	possibleDoubleCountedSerials        []string
+	possibleDoubleCountedRevokedSerials []string
+
 	pkiStorageVersion atomic.Value
 	crlBuilder        *crlBuilder
 
@@ -325,6 +341,21 @@ func (b *backend) initialize(ctx context.Context, _ *logical.InitializationReque
 		return err
 	}
 
+	err := b.initializePKIIssuersStorage(ctx)
+	if err != nil {
+		return err
+	}
+
+	// Initialize also needs to populate our certificate and revoked certificate count
+	err = b.initializeStoredCertificateCounts(ctx)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (b *backend) initializePKIIssuersStorage(ctx context.Context) error {
 	// Grab the lock prior to the updating of the storage lock preventing us flipping
 	// the storage flag midway through the request stream of other requests.
 	b.issuersLock.Lock()
@@ -527,3 +558,154 @@ func (b *backend) periodicFunc(ctx context.Context, request *logical.Request) er
 	// All good!
 	return nil
 }
+
+func (b *backend) initializeStoredCertificateCounts(ctx context.Context) error {
+	b.tidyStatusLock.RLock()
+	defer b.tidyStatusLock.RUnlock()
+	// For performance reasons, we can't lock on issuance/storage of certs until a list operation completes,
+	// but we want to limit possible miscounts / double-counts to over-counting, so we take the tidy lock which
+	// prevents (most) deletions - in particular we take a read lock (sufficient to block the write lock in
+	// tidyStatusStart while allowing tidy to still acquire a read lock to report via its endpoint)
+
+	entries, err := b.storage.List(ctx, "certs/")
+	if err != nil {
+		return err
+	}
+	atomic.AddUint32(b.certCount, uint32(len(entries)))
+
+	revokedEntries, err := b.storage.List(ctx, "revoked/")
+	if err != nil {
+		return err
+	}
+	atomic.AddUint32(b.revokedCertCount, uint32(len(revokedEntries)))
+
+	b.certsCounted.Store(true)
+	// Now that the metrics are set, we can switch from appending newly-stored certificates to the possible double-count
+	// list, and instead have them update the counter directly.  We need to do this so that we are looking at a static
+	// slice of possibly double counted serials.  Note that certsCounted is computed before the storage operation, so
+	// there may be some delay here.
+
+	// Sort the listed-entries first, to accommodate that delay.
+	sort.Slice(entries, func(i, j int) bool {
+		return entries[i] < entries[j]
+	})
+
+	sort.Slice(revokedEntries, func(i, j int) bool {
+		return revokedEntries[i] < revokedEntries[j]
+	})
+
+	// We assume here that these lists are now complete.
+	sort.Slice(b.possibleDoubleCountedSerials, func(i, j int) bool {
+		return b.possibleDoubleCountedSerials[i] < b.possibleDoubleCountedSerials[j]
+	})
+
+	listEntriesIndex := 0
+	possibleDoubleCountIndex := 0
+	for {
+		if listEntriesIndex >= len(entries) {
+			break
+		}
+		if possibleDoubleCountIndex >= len(b.possibleDoubleCountedSerials) {
+			break
+		}
+		if entries[listEntriesIndex] == b.possibleDoubleCountedSerials[possibleDoubleCountIndex][6:] {
+			// This represents a double-counted entry
+			b.decrementTotalCertificatesCountNoReport()
+			listEntriesIndex = listEntriesIndex + 1
+			possibleDoubleCountIndex = possibleDoubleCountIndex + 1
+			continue
+		}
+		if entries[listEntriesIndex] < b.possibleDoubleCountedSerials[possibleDoubleCountIndex][6:] {
+			listEntriesIndex = listEntriesIndex + 1
+			continue
+		}
+		if entries[listEntriesIndex] > b.possibleDoubleCountedSerials[possibleDoubleCountIndex] {
+			possibleDoubleCountIndex = possibleDoubleCountIndex + 1
+			continue
+		}
+	}
+
+	sort.Slice(b.possibleDoubleCountedRevokedSerials, func(i, j int) bool {
+		return b.possibleDoubleCountedRevokedSerials[i] < b.possibleDoubleCountedRevokedSerials[j]
+	})
+
+	listRevokedEntriesIndex := 0
+	possibleRevokedDoubleCountIndex := 0
+	for {
+		if listRevokedEntriesIndex >= len(revokedEntries) {
+			break
+		}
+		if possibleRevokedDoubleCountIndex >= len(b.possibleDoubleCountedRevokedSerials) {
+			break
+		}
+		if revokedEntries[listRevokedEntriesIndex] == b.possibleDoubleCountedRevokedSerials[possibleRevokedDoubleCountIndex][8:] {
+			// This represents a double-counted revoked entry
+			b.decrementTotalRevokedCertificatesCountNoReport()
+			listRevokedEntriesIndex = listRevokedEntriesIndex + 1
+			possibleRevokedDoubleCountIndex = possibleRevokedDoubleCountIndex + 1
+			continue
+		}
+		if revokedEntries[listRevokedEntriesIndex] < b.possibleDoubleCountedRevokedSerials[possibleRevokedDoubleCountIndex][8:] {
+			listRevokedEntriesIndex = listRevokedEntriesIndex + 1
+			continue
+		}
+		if revokedEntries[listRevokedEntriesIndex] > b.possibleDoubleCountedRevokedSerials[possibleRevokedDoubleCountIndex][8:] {
+			possibleRevokedDoubleCountIndex = possibleRevokedDoubleCountIndex + 1
+			continue
+		}
+	}
+
+	b.possibleDoubleCountedRevokedSerials = nil
+	b.possibleDoubleCountedSerials = nil
+
+	metrics.SetGauge([]string{"secrets", "pki", "total_certificates_stored"}, float32(*b.certCount))
+	metrics.SetGauge([]string{"secrets", "pki", "total_revoked_certificates_stored"}, float32(*b.revokedCertCount))
+
+	return nil
+}
+
+// The "certsCounted" boolean here should be loaded from the backend certsCounted before the corresponding storage call:
+// eg. certsCounted := b.certsCounted.Load()
+func (b *backend) incrementTotalCertificatesCount(certsCounted bool, newSerial string) {
+	atomic.AddUint32(b.certCount, 1)
+	switch {
+	case !certsCounted:
+		// This is unsafe, but a good best-attempt
+		b.possibleDoubleCountedSerials = append(b.possibleDoubleCountedSerials, newSerial)
+	default:
+		metrics.SetGauge([]string{"secrets", "pki", "total_certificates_stored"}, float32(*b.certCount))
+	}
+}
+
+func (b *backend) decrementTotalCertificatesCountReport() {
+	b.decrementTotalCertificatesCountNoReport()
+	metrics.SetGauge([]string{"secrets", "pki", "total_certificates_stored"}, float32(*b.certCount))
+}
+
+// Called directly only by the initialize function to deduplicate the count, when we don't have a full count yet
+func (b *backend) decrementTotalCertificatesCountNoReport() {
+	atomic.AddUint32(b.certCount, ^uint32(0))
+}
+
+// The "certsCounted" boolean here should be loaded from the backend certsCounted before the corresponding storage call:
+// eg. certsCounted := b.certsCounted.Load()
+func (b *backend) incrementTotalRevokedCertificatesCount(certsCounted bool, newSerial string) {
+	atomic.AddUint32(b.revokedCertCount, 1)
+	switch {
+	case !certsCounted:
+		// This is unsafe, but a good best-attempt
+		b.possibleDoubleCountedRevokedSerials = append(b.possibleDoubleCountedRevokedSerials, newSerial)
+	default:
+		metrics.SetGauge([]string{"secrets", "pki", "total_revoked_certificates_stored"}, float32(*b.revokedCertCount))
+	}
+}
+
+func (b *backend) decrementTotalRevokedCertificatesCountReport() {
+	b.decrementTotalRevokedCertificatesCountNoReport()
+	metrics.SetGauge([]string{"secrets", "pki", "total_revoked_certificates_stored"}, float32(*b.revokedCertCount))
+}
+
+// Called directly only by the initialize function to deduplicate the count, when we don't have a full count yet
+func (b *backend) decrementTotalRevokedCertificatesCountNoReport() {
+	atomic.AddUint32(b.revokedCertCount, ^uint32(0))
+}