fix(treatment-service): remove nested RLock deadlock and add optional LocalStorage call metrics

plugins/turing/runner/experiment_runner.go

@@ -238,6 +238,19 @@ func (er *experimentRunner) RegisterMetricsCollector(
 				instrumentation.AdditionalNoMatchingExperimentRequestCountLabels...,
 			),
 		},
+		{
+			Name:        string(instrumentation.LocalStorageCallCount),
+			Type:        routerMetrics.CounterMetricType,
+			Description: instrumentation.LocalStorageCallCountHelpString,
+			Labels:      instrumentation.LocalStorageMethodLabels,
+		},
+		{
+			Name:        string(instrumentation.LocalStorageCallDurationMs),
+			Type:        routerMetrics.HistogramMetricType,
+			Description: instrumentation.LocalStorageCallDurationMsHelpString,
+			Buckets:     instrumentation.RequestLatencyBuckets,
+			Labels:      instrumentation.LocalStorageMethodLabels,
+		},
 	})
 	if err != nil {
 		return err

treatment-service/appcontext/appcontext.go

@@ -67,6 +67,9 @@ func NewAppContext(cfg *config.Config) (*AppContext, error) {
 	if err != nil {
 		return nil, err
 	}
+	if cfg.MonitoringConfig.LocalStorageMetricsEnabled {
+		localStorage.SetMetricsRecorder(metricService)
+	}
 
 	log.Println("Initializing assigned treatment logger...")
 	loggerConfig := cfg.AssignedTreatmentLogger

treatment-service/config/config.go

@@ -87,6 +87,10 @@ const (
 type Monitoring struct {
 	Kind         MetricSinkKind `json:"kind" default:"" validate:"required"`
 	MetricLabels []string       `json:"metric_labels" default:""`
+	// LocalStorageMetricsEnabled toggles per-method call counter and call duration
+	// histogram instrumentation on models.LocalStorage's public methods. Independent of
+	// Kind, since it self-registers directly against the default Prometheus registry.
+	LocalStorageMetricsEnabled bool `json:"local_storage_metrics_enabled" default:"false"`
 }
 
 type ManagementServiceConfig struct {

treatment-service/instrumentation/prometheus.go

@@ -26,6 +26,14 @@ const (
 	FetchTreatmentRequestCountHelpString string = "Counter for no. of Fetch Treatment requests with matching experiments"
 	// NoMatchingExperimentRequestCountHelpString is the help string of the NoMatchingExperimentRequestCount metric
 	NoMatchingExperimentRequestCountHelpString string = "Counter for no. of Fetch Treatment requests with no matching experiments"
+	// LocalStorageCallCount is the key to measure calls to LocalStorage's public methods
+	LocalStorageCallCount metrics.MetricName = "local_storage_calls_total"
+	// LocalStorageCallDurationMs is the key to measure how long a call to a LocalStorage public method took
+	LocalStorageCallDurationMs metrics.MetricName = "local_storage_call_duration_ms"
+	// LocalStorageCallCountHelpString is the help string of the LocalStorageCallCount metric
+	LocalStorageCallCountHelpString string = "Counter for calls to LocalStorage's public methods, incremented on entry"
+	// LocalStorageCallDurationMsHelpString is the help string of the LocalStorageCallDurationMs metric
+	LocalStorageCallDurationMsHelpString string = "Histogram for how long (in milliseconds) a call to a LocalStorage public method took to return"
 )
 
 // RequestLatencyBuckets defines the buckets used in the custom Histogram metrics
@@ -49,6 +57,12 @@ var FetchTreatmentRequestDurationMsLabels = []string{"project_name", "experiment
 // ExperimentLookupDurationMsLabels defines additional labels needed for the ExperimentLookupDurationMs histogram map
 var ExperimentLookupDurationMsLabels = []string{"project_name"}
 
+// LocalStorageMethodLabels defines the labels needed for the LocalStorageCallCount counter map and the
+// LocalStorageCallDurationMs histogram map. Deliberately not combined with the caller-supplied custom
+// MetricLabels (cfg.MetricLabels) used elsewhere in this file -- LocalStorage calls aren't tied to a project
+// or experiment, just a method name.
+var LocalStorageMethodLabels = []string{"method"}
+
 var GaugeMap = map[metrics.MetricName]metrics.PrometheusGaugeVec{}
 
 func GetCounterMap(labels []string) map[metrics.MetricName]metrics.PrometheusCounterVec {
@@ -76,6 +90,14 @@ func GetCounterMap(labels []string) map[metrics.MetricName]metrics.PrometheusCou
 		},
 			noMatchingExperimentlabels,
 		),
+		LocalStorageCallCount: prometheus.NewCounterVec(prometheus.CounterOpts{
+			Namespace: Namespace,
+			Subsystem: Subsystem,
+			Help:      LocalStorageCallCountHelpString,
+			Name:      string(LocalStorageCallCount),
+		},
+			LocalStorageMethodLabels,
+		),
 	}
 
 	return counterMap
@@ -102,6 +124,15 @@ func GetHistogramMap() map[metrics.MetricName]metrics.PrometheusHistogramVec {
 		},
 			ExperimentLookupDurationMsLabels,
 		),
+		LocalStorageCallDurationMs: prometheus.NewHistogramVec(prometheus.HistogramOpts{
+			Namespace: Namespace,
+			Subsystem: Subsystem,
+			Name:      string(LocalStorageCallDurationMs),
+			Help:      LocalStorageCallDurationMsHelpString,
+			Buckets:   RequestLatencyBuckets,
+		},
+			LocalStorageMethodLabels,
+		),
 	}
 
 	return histogramMap

treatment-service/models/storage.go

@@ -59,6 +59,28 @@ type LocalStorage struct {
 	subscribedProjectIds []ProjectId
 	Segmenters           map[string]schema.SegmenterType
 	ProjectSegmenters    map[ProjectId]map[string]schema.SegmenterType
+	// metricsRecorder is nil until SetMetricsRecorder is called; instrumented methods treat a
+	// nil recorder as "metrics disabled" and skip reporting entirely. In production this is a
+	// services.MetricService (wired in appcontext.NewAppContext, after both are constructed --
+	// models can't import services directly, since services already imports models).
+	//
+	// Every instrumented public method here calls into metricsRecorder.LogRequestCount /
+	// LogLatencyHistogram via instrumentCall. services.MetricService itself already calls back
+	// into this struct: GetLabels and GetProjectNameLabel both call FindProjectSettingsWithId,
+	// which is one of the instrumented methods. That's fine today because LogRequestCount and
+	// LogLatencyHistogram are pure leaves -- they only reach the external metrics collector and
+	// never call back into GetLabels, GetProjectNameLabel, or any LocalStorage method. If you
+	// ever change what a recorder's LogRequestCount/LogLatencyHistogram does, make sure it still
+	// can't reach back into an instrumented LocalStorage method (directly or transitively) --
+	// that would turn this into unbounded recursion, not just a slow call.
+	metricsRecorder LocalStorageMetricsRecorder
+}
+
+// SetMetricsRecorder wires recorder into LocalStorage's public methods, which will start
+// reporting a call counter and call duration histogram to it. Call once, after construction;
+// leaving it unset (the default) means no metrics are reported.
+func (s *LocalStorage) SetMetricsRecorder(recorder LocalStorageMetricsRecorder) {
+	s.metricsRecorder = recorder
 }
 
 type Match struct {
@@ -254,6 +276,10 @@ func (i *ExperimentIndex) checkSegmentHasWeakMatch(segmentName string) bool {
 }
 
 func (s *LocalStorage) InsertProjectSettings(projectSettings *pubsub.ProjectSettings) error {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("InsertProjectSettings")()
+	}
+
 	// check that settings with the same Id doesn't exist
 	existingProjectSettings := s.findProjectSettingsById(ProjectId(projectSettings.GetProjectId()))
 	if existingProjectSettings != nil {
@@ -274,6 +300,10 @@ func (s *LocalStorage) InsertProjectSettings(projectSettings *pubsub.ProjectSett
 }
 
 func (s *LocalStorage) UpdateProjectSettings(updatedProjectSettings *pubsub.ProjectSettings) {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("UpdateProjectSettings")()
+	}
+
 	s.Lock()
 	defer s.Unlock()
 
@@ -285,6 +315,10 @@ func (s *LocalStorage) UpdateProjectSettings(updatedProjectSettings *pubsub.Proj
 }
 
 func (s *LocalStorage) FindProjectSettingsWithId(projectId ProjectId) *pubsub.ProjectSettings {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("FindProjectSettingsWithId")()
+	}
+
 	projectSettings := s.findSubscribedProjectSettingsById(projectId)
 	if projectSettings != nil {
 		return projectSettings
@@ -307,13 +341,21 @@ func (s *LocalStorage) findSubscribedProjectSettingsById(projectId ProjectId) *p
 		return nil
 	}
 
-	return s.findProjectSettingsById(projectId)
+	return s.findProjectSettingsByIdLocked(projectId)
 }
 
 func (s *LocalStorage) findProjectSettingsById(projectId ProjectId) *pubsub.ProjectSettings {
 	s.RLock()
 	defer s.RUnlock()
 
+	return s.findProjectSettingsByIdLocked(projectId)
+}
+
+// findProjectSettingsByIdLocked assumes the caller already holds s's read (or write) lock.
+// It must not itself call RLock/Lock, since sync.RWMutex blocks new readers once a writer is
+// queued -- a nested RLock on the same goroutine that already holds the lock would deadlock
+// against that queued writer.
+func (s *LocalStorage) findProjectSettingsByIdLocked(projectId ProjectId) *pubsub.ProjectSettings {
 	for _, settings := range s.ProjectSettings {
 		if ProjectId(settings.ProjectId) == projectId {
 			return settings
@@ -337,6 +379,10 @@ func (s *LocalStorage) fetchProjectSettingsWithId(projectId ProjectId) (*pubsub.
 }
 
 func (s *LocalStorage) GetSegmentersTypeMapping(projectId ProjectId) (map[string]schema.SegmenterType, error) {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("GetSegmentersTypeMapping")()
+	}
+
 	s.RLock()
 	defer s.RUnlock()
 
@@ -348,6 +394,10 @@ func (s *LocalStorage) GetSegmentersTypeMapping(projectId ProjectId) (map[string
 }
 
 func (s *LocalStorage) FindExperiments(projectId ProjectId, filters []SegmentFilter) []*ExperimentMatch {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("FindExperiments")()
+	}
+
 	s.RLock()
 	defer s.RUnlock()
 
@@ -382,6 +432,10 @@ func (s *LocalStorage) FindExperiments(projectId ProjectId, filters []SegmentFil
 }
 
 func (s *LocalStorage) FindExperimentWithId(projectId ProjectId, experimentId int64) *pubsub.Experiment {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("FindExperimentWithId")()
+	}
+
 	s.RLock()
 	defer s.RUnlock()
 
@@ -454,6 +508,10 @@ func NewExperimentIndex(experiment *pubsub.Experiment) *ExperimentIndex {
 }
 
 func (s *LocalStorage) InsertExperiment(experiment *pubsub.Experiment) {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("InsertExperiment")()
+	}
+
 	projectId := ProjectId(experiment.ProjectId)
 	s.Lock()
 	defer s.Unlock()
@@ -475,6 +533,10 @@ func (s *LocalStorage) InsertExperiment(experiment *pubsub.Experiment) {
 }
 
 func (s *LocalStorage) UpdateExperiment(experiment *pubsub.Experiment) {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("UpdateExperiment")()
+	}
+
 	projectId := ProjectId(experiment.ProjectId)
 	s.Lock()
 	defer s.Unlock()
@@ -502,6 +564,10 @@ func (s *LocalStorage) UpdateExperiment(experiment *pubsub.Experiment) {
 // DumpExperiments is used to dump the experiment from the local cache into the
 // given file, as JSON. Useful for debugging.
 func (s *LocalStorage) DumpExperiments(filepath string) error {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("DumpExperiments")()
+	}
+
 	s.RLock()
 	defer s.RUnlock()
 
@@ -513,6 +579,10 @@ func (s *LocalStorage) DumpExperiments(filepath string) error {
 }
 
 func (s *LocalStorage) Init() error {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("Init")()
+	}
+
 	var subscribedProjectSettings []*pubsub.ProjectSettings
 	var err error
 	if len(s.subscribedProjectIds) > 0 {
@@ -621,7 +691,11 @@ func NewLocalStorage(
 		return nil, err
 	}
 	segmenterCache := make(map[ProjectId]map[string]schema.SegmenterType)
-	s := LocalStorage{managementClient: xpClient, subscribedProjectIds: projectIds, ProjectSegmenters: segmenterCache}
+	s := LocalStorage{
+		managementClient:     xpClient,
+		subscribedProjectIds: projectIds,
+		ProjectSegmenters:    segmenterCache,
+	}
 	err = s.Init()
 
 	return &s, err
@@ -709,12 +783,20 @@ func (s *LocalStorage) fetchProjectSegmenters(settings []*pubsub.ProjectSettings
 }
 
 func (s *LocalStorage) UpdateProjectSegmenters(segmenter *_segmenters.SegmenterConfiguration, projectId int64) {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("UpdateProjectSegmenters")()
+	}
+
 	s.Lock()
 	defer s.Unlock()
 	s.ProjectSegmenters[ProjectId(projectId)][segmenter.Name] = schema.SegmenterType(strings.ToLower(segmenter.Type.String()))
 }
 
 func (s *LocalStorage) DeleteProjectSegmenters(segmenterName string, projectId int64) {
+	if s.metricsRecorder != nil {
+		defer s.instrumentCall("DeleteProjectSegmenters")()
+	}
+
 	s.Lock()
 	defer s.Unlock()
 	delete(s.ProjectSegmenters[ProjectId(projectId)], segmenterName)