[m3coordinator] Allow downsampler to configure custom rule store (#2996)

src/cmd/services/m3coordinator/downsample/options.go

@@ -102,6 +102,9 @@ var (
 	errRollupRuleNoTransforms       = errors.New("rollup rule has no transforms set")
 )
 
+// CustomRuleStoreFn is a function to swap the backend used for the rule stores.
+type CustomRuleStoreFn func(kv.Store) (kv.Store, error)
+
 // DownsamplerOptions is a set of required downsampler options.
 type DownsamplerOptions struct {
 	Storage                    storage.Storage

src/query/server/query.go

@@ -149,6 +149,9 @@ type RunOptions struct {
 	// CustomBuildTags are additional tags to be added to the instrument build
 	// reporter.
 	CustomBuildTags map[string]string
+
+	// ApplyCustomRuleStore provides an option to swap the backend used for the rule stores.
+	ApplyCustomRuleStore downsample.CustomRuleStoreFn
 }
 
 // InstrumentOptionsReady is a set of instrument options
@@ -669,7 +672,7 @@ func newM3DBStorage(
 		ds, err := newDownsampler(
 			cfg.Downsample, clusterClient,
 			fanoutStorage, clusterNamespacesWatcher,
-			tsdbOpts.TagOptions(), instrumentOptions, rwOpts)
+			tsdbOpts.TagOptions(), instrumentOptions, rwOpts, runOpts.ApplyCustomRuleStore)
 		if err != nil {
 			return nil, err
 		}
@@ -728,6 +731,7 @@ func newDownsampler(
 	tagOptions models.TagOptions,
 	instrumentOpts instrument.Options,
 	rwOpts xio.Options,
+	applyCustomRuleStore downsample.CustomRuleStoreFn,
 ) (downsample.Downsampler, error) {
 	// Namespace the downsampler metrics.
 	instrumentOpts = instrumentOpts.SetMetricsScope(
@@ -743,6 +747,12 @@ func newDownsampler(
 		return nil, errors.Wrap(err, "unable to create KV store from the "+
 			"cluster management config client")
 	}
+	if applyCustomRuleStore != nil {
+		kvStore, err = applyCustomRuleStore(kvStore)
+		if err != nil {
+			return nil, errors.Wrap(err, "unable to apply custom rule store")
+		}
+	}
 
 	tagEncoderOptions := serialize.NewTagEncoderOptions()
 	tagDecoderOptions := serialize.NewTagDecoderOptions(serialize.TagDecoderOptionsConfig{})