Add mean as an additional way to collapse probabilities

pkg/config/tuning.go

@@ -1,6 +1,7 @@
 package config
 
 import (
+	"log"
 	"math"
 	"time"
 )
@@ -27,6 +28,37 @@ const (
 	defaultRotationDuration = time.Minute * 5
 )
 
+// The function to choose the final probability based on all bucket probabilities
+type FinalProbabilityFunction func([]float64) float64
+
+var (
+	MinFinalProbabilityFunction FinalProbabilityFunction = func(buckets []float64) float64 {
+		if len(buckets) == 0 {
+			log.Fatalf("Cannot compute final probability with empty buckets slice")
+		}
+
+		var min float64 = 1.
+		for _, b := range buckets {
+			min = math.Min(min, b)
+		}
+
+		return min
+	}
+
+	MeanFinalProbabilityFunction FinalProbabilityFunction = func(buckets []float64) float64 {
+		if len(buckets) == 0 {
+			log.Fatalf("Cannot compute final probability with empty buckets slice")
+		}
+
+		var total float64
+		for _, b := range buckets {
+			total += b
+		}
+
+		return total / float64(len(buckets))
+	}
+)
+
 // The default config that's supposed to work in most cases
 func DefaultFairnessTrackerConfig() *FairnessTrackerConfig {
 	return GenerateTunedStructureConfig(
@@ -58,12 +90,14 @@ func GenerateTunedStructureConfig(expectedClientFlows, bucketsPerLevel, tolerabl
 	Pd := pdSlowingFactor * Pi
 
 	return &FairnessTrackerConfig{
-		M:                 bucketsPerLevel,
-		L:                 L,
-		Pi:                Pi,
-		Pd:                Pd,
-		Lambda:            defaultDecayRate,
-		RotationFrequency: defaultRotationDuration,
+		M:                        bucketsPerLevel,
+		L:                        L,
+		Pi:                       Pi,
+		Pd:                       Pd,
+		Lambda:                   defaultDecayRate,
+		RotationFrequency:        defaultRotationDuration,
+		IncludeStats:             false,
+		FinalProbabilityFunction: MinFinalProbabilityFunction,
 	}
 }
 

pkg/config/types.go

@@ -18,4 +18,6 @@ type FairnessTrackerConfig struct {
 	RotationFrequency time.Duration
 	// Include result stats. Useful for debugging but may slightly affect performance.
 	IncludeStats bool
+	// The function to choose the final probability from all the bucket probabilities
+	FinalProbabilityFunction FinalProbabilityFunction
 }

pkg/data/data.go

@@ -89,33 +89,32 @@ func (s *Structure) Close() {
 func (s *Structure) RegisterRequest(ctx context.Context, clientIdentifier []byte) (*request.RegisterRequestResult, error) {
 	var stats *request.ResultStats
 
-	var pmin float64 = 1
+	bucketProbabilities := make([]float64, s.config.L)
 
 	// We can ignore the error since the handler never returns one
 	s.visitBuckets(clientIdentifier, func(l uint32, m uint32, b *bucket) error {
-		if b.probability < pmin {
-			pmin = b.probability
-		}
-
+		bucketProbabilities[l] = b.probability
 		if s.includeStats {
 			if stats == nil {
 				stats = &request.ResultStats{
-					BucketIndexes:       make([]int, s.config.L),
-					BucketProbabilities: make([]float64, s.config.L),
+					BucketIndexes: make([]int, s.config.L),
 				}
 			}
-
 			stats.BucketIndexes[l] = int(m)
-			stats.BucketProbabilities[l] = b.probability
-			stats.FinalProbability = pmin
 		}
-
 		return nil
 	})
 
+	pfinal := s.config.FinalProbabilityFunction(bucketProbabilities)
+
+	if s.includeStats {
+		stats.BucketProbabilities = bucketProbabilities
+		stats.FinalProbability = pfinal
+	}
+
 	// Decide whether to throttle the request based on the probability
 	shouldThrottle := false
-	if rand.Float64() <= pmin {
+	if rand.Float64() <= pfinal {
 		shouldThrottle = true
 	}
 

pkg/data/data_test.go

@@ -121,10 +121,12 @@ func TestGetId(t *testing.T) {
 
 func TestEndToEnd(t *testing.T) {
 	conf := &config.FairnessTrackerConfig{
-		L:  2,
-		M:  24,
-		Pd: .1,
-		Pi: .15,
+		L:                        2,
+		M:                        24,
+		Pd:                       .1,
+		Pi:                       .15,
+		Lambda:                   0,
+		FinalProbabilityFunction: config.MeanFinalProbabilityFunction,
 	}
 	structure, err := NewStructure(conf, 1, true)
 	assert.NoError(t, err)

pkg/tracker/types.go

@@ -58,6 +58,10 @@ func (bl *FairnessTrackerBuilder) SetRotationFrequency(rotationFrequency time.Du
 	bl.configuration.RotationFrequency = rotationFrequency
 }
 
+func (bl *FairnessTrackerBuilder) SetFinalProbabilityFunction(finalProbabilityFunction config.FinalProbabilityFunction) {
+	bl.configuration.FinalProbabilityFunction = finalProbabilityFunction
+}
+
 // The public facing errors from the FairnessTracker
 type FairnessTrackerError struct {
 	*utils.BaseError

pkg/tracker/types_test.go

@@ -26,6 +26,7 @@ func TestBuildFairnessTracker(t *testing.T) {
 	b.SetLambda(.001)
 	b.SetRotationFrequency(1 * time.Second)
 	b.SetIncludeStats(true)
+	b.SetFinalProbabilityFunction(config.MeanFinalProbabilityFunction)
 
 	tr, err := b.Build()
 	assert.NoError(t, err)