From c2d30e2ffc632a1ea64030467e5a40e02e4158be Mon Sep 17 00:00:00 2001
From: Ivan Kozlovic <ivan@synadia.com>
Date: Thu, 22 Feb 2024 09:57:27 -0700
Subject: [PATCH] [ADDED] Sampling in MsgTrace structure (#219)

When an account has a `Trace` field set with a `Destination`, if
an incoming message has a `traceparent` header with proper flag
indicating that the message should be trace, the `Sampling` integer
(value in the range of [1..100]) allows tracing only a random
sample of messages. The server will pick a random number below 100
and if that number is lower or equal to this `Sampling` value
then the trace will trigger. If above, the trace will not trigger.

Signed-off-by: Ivan Kozlovic <ivan@synadia.com>
---
 v2/account_claims.go      | 14 ++++++++++++
 v2/account_claims_test.go | 46 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 60 insertions(+)

diff --git a/v2/account_claims.go b/v2/account_claims.go
index 193516c..5c1665d 100644
--- a/v2/account_claims.go
+++ b/v2/account_claims.go
@@ -237,7 +237,16 @@ type Account struct {
 
 // MsgTrace holds distributed message tracing configuration
 type MsgTrace struct {
+	// Destination is the subject the server will send message traces to
+	// if the inbound message contains the "traceparent" header and has
+	// its sampled field indicating that the trace should be triggered.
 	Destination Subject `json:"dest,omitempty"`
+	// Sampling is used to set the probability sampling, that is, the
+	// server will get a random number between 1 and 100 and trigger
+	// the trace if the number is lower than this Sampling value.
+	// The valid range is [1..100]. If the value is not set Validate()
+	// will set the value to 100.
+	Sampling int `json:"sampling,omitempty"`
 }
 
 // Validate checks if the account is valid, based on the wrapper
@@ -257,6 +266,11 @@ func (a *Account) Validate(acct *AccountClaims, vr *ValidationResults) {
 		if a.Trace.Destination.HasWildCards() {
 			vr.AddError("the account Trace.Destination subject %q is not a valid publish subject", a.Trace.Destination)
 		}
+		if a.Trace.Sampling < 0 || a.Trace.Sampling > 100 {
+			vr.AddError("the account Trace.Sampling value '%d' is not valid, should be in the range [1..100]", a.Trace.Sampling)
+		} else if a.Trace.Sampling == 0 {
+			a.Trace.Sampling = 100
+		}
 	}
 
 	if !a.Limits.IsEmpty() && a.Limits.Imports >= 0 && int64(len(a.Imports)) > a.Limits.Imports {
diff --git a/v2/account_claims_test.go b/v2/account_claims_test.go
index dd5a27e..19dbf93 100644
--- a/v2/account_claims_test.go
+++ b/v2/account_claims_test.go
@@ -17,6 +17,7 @@ package jwt
 
 import (
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -910,3 +911,48 @@ func TestAccountClaimsTraceDest(t *testing.T) {
 		})
 	}
 }
+
+func TestAccountClaimsTraceDestSampling(t *testing.T) {
+	akp := createAccountNKey(t)
+	apk := publicKey(akp, t)
+
+	account := NewAccountClaims(apk)
+	for _, test := range []struct {
+		name      string
+		dest      string
+		sampling  int
+		expectErr string
+	}{
+		{"sampling without destination", "", 10, "subject cannot be empty"},
+		{"sampling negative", "dest", -1, "should be in the range [1..100]"},
+		{"sampling above 100", "dest", 101, "should be in the range [1..100]"},
+		{"sampling at 50", "dest", 50, ""},
+		{"sampling at zero sets to 100", "dest", 0, ""},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			account.Trace = &MsgTrace{Destination: Subject(test.dest), Sampling: test.sampling}
+			vr := CreateValidationResults()
+			account.Validate(vr)
+
+			if test.expectErr == "" {
+				if !vr.IsEmpty() {
+					t.Fatalf("account validation should not have failed, got %+v", vr.Issues)
+				}
+				if test.sampling == 0 {
+					if account.Trace.Sampling != 100 {
+						t.Fatalf("account sampling should have been set to 100 got %d", account.Trace.Sampling)
+					}
+				} else if test.sampling != account.Trace.Sampling {
+					t.Fatalf("account sampling should be %d, got %d", test.sampling, account.Trace.Sampling)
+				}
+			} else {
+				if vr.IsEmpty() {
+					t.Fatal("account validation should have failed")
+				}
+				if !strings.Contains(vr.Issues[0].Description, test.expectErr) {
+					t.Fatalf("account validation should have failed with error %q, got %q", test.expectErr, vr.Issues[0].Description)
+				}
+			}
+		})
+	}
+}