Metrics: Add auto tracked files mode

When enabled, files will be automatically added to tracked files.

Author: Skantes

cli/commands.go

@@ -355,6 +355,9 @@ func (c *cli) CmdMetrics(args ...string) error {
 	_ = cmd.String("add", "", "Add a file to the metrics route")
 	_ = cmd.String("list", "*", "List files in metrics + Optionnal pattern to filter results")
 	_ = cmd.String("delete", "", "Delete a file from the metrics route")
+	_ = cmd.Bool("auto-enable", true, "Enable automatic addition of new files to tracked files")
+	_ = cmd.Bool("auto-disable", false, "Disable automatic addition of new files to tracked files")
+	_ = cmd.Bool("auto-status", true, "Print boolean of automatic addition of new files to tracked files")
 
 	// Can't use cmd.Parse(args) because it doesn't handle -command without
 	// an argument following which is needed for list
@@ -372,6 +375,12 @@ func (c *cli) CmdMetrics(args ...string) error {
 		c.CmdAddmetric(args[1])
 	} else if args[0] == "-delete" && len(args) == 2 {
 		c.CmdDelmetric(args[1])
+	} else if args[0] == "-auto-enable" {
+		c.CmdEnableauto()
+	} else if args[0] == "-auto-disable" {
+		c.CmdDisableauto()
+	} else if args[0] == "-auto-status" {
+		c.CmdStatusauto()
 	} else {
 		cmd.Usage()
 		return nil
@@ -442,6 +451,37 @@ func (c *cli) CmdListmetrics(pattern string) error {
 	return nil
 }
 
+func (c *cli) CmdEnableauto() error {
+	client := c.GetRPC()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultRPCTimeout)
+	defer cancel()
+	client.EnableAuto(ctx, &empty.Empty{})
+	return nil
+}
+
+func (c *cli) CmdDisableauto() error {
+	client := c.GetRPC()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultRPCTimeout)
+	defer cancel()
+	client.DisableAuto(ctx, &empty.Empty{})
+	return nil
+}
+
+func (c *cli) CmdStatusauto() error {
+	client := c.GetRPC()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultRPCTimeout)
+	defer cancel()
+	status, _ := client.GetStatusAuto(ctx, &empty.Empty{})
+
+	if status.Status {
+		log.Info("Auto tracked files is enabled")
+	} else {
+		log.Info("Auto tracked files is disabled")
+	}
+
+	return nil
+}
+
 func (c *cli) CmdRemove(args ...string) error {
 	cmd := SubCmd("remove", "IDENTIFIER", "Remove an existing mirror")
 	force := cmd.Bool("f", false, "Never prompt for confirmation")

config/config.go

@@ -61,6 +61,7 @@ func defaultConfig() Configuration {
 		RPCPassword:              "",
 		MetricsEnabled:           false,
 		MetricsTopFilesRetention: 0,
+		MetricsAutoTrackedFiles:  false,
 	}
 }
 
@@ -98,6 +99,7 @@ type Configuration struct {
 
 	MetricsEnabled           bool `yaml:"MetricsEnabled"`
 	MetricsTopFilesRetention int  `yaml:"MetricsTopFilesRetention"`
+	MetricsAutoTrackedFiles  bool `yaml:"MetricsAutoTrackedFiles"`
 }
 
 type fallback struct {

database/utils.go

@@ -4,10 +4,12 @@
 package database
 
 import (
-	"errors"
 	"strconv"
 
 	"github.com/gomodule/redigo/redis"
+	"github.com/pkg/errors"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func (r *Redis) GetListOfMirrors() (map[int]string, error) {
@@ -147,6 +149,52 @@ func (r *Redis) AddCountry(country string) error {
 	return err
 }
 
+func (r *Redis) AddTrackedFile(file string) error {
+	conn, err := r.Connect()
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	exists, err := redis.Int(conn.Do("SISMEMBER", "FILES", file))
+	if err != nil {
+		return errors.Wrap(err, "failed to check for file presence")
+	} else if exists == 0 {
+		return status.Error(codes.FailedPrecondition,
+			"file does not exist")
+	}
+
+	exists, err = redis.Int(conn.Do("SADD", "TRACKED_FILES", file))
+	if err != nil {
+		return errors.Wrap(err, "failed to add file to metrics")
+	} else if exists == 0 {
+		return status.Error(codes.AlreadyExists, "file already is in metrics")
+	}
+	return nil
+}
+
+func (r *Redis) DeleteTrackedFile(file string) error {
+	conn, err := r.Connect()
+	if err != nil {
+		return err
+	}
+	defer conn.Close()
+
+	exists, err := redis.Int(conn.Do("SISMEMBER", "TRACKED_FILES", file))
+	if err != nil {
+		return errors.Wrap(err, "failed to check for file presence")
+	} else if exists == 0 {
+		return status.Error(codes.FailedPrecondition,
+			"file is not part of tracked files")
+	}
+
+	_, err = conn.Do("SREM", "TRACKED_FILES", file)
+	if err != nil {
+		return errors.Wrap(err, "failed to remove file from tracked files")
+	}
+	return nil
+}
+
 type NetReadyError struct {
 	error
 }

mirrorbits.conf

@@ -140,3 +140,6 @@
 ## Keep in mind that this will have an impact on the database's RAM usage
 ## It is not recommended to keep more than a few days of retention.
 ## MetricsTopFilesRetention: 0
+
+## Enable / Disable automatically adding a new file to tracked files
+## MetricsAutoTrackedFiles: false

rpc/rpc.go

@@ -730,52 +730,11 @@ func (c *CLI) GetMirrorLogs(ctx context.Context, in *GetMirrorLogsRequest) (*Get
 }
 
 func (c *CLI) AddMetric(ctx context.Context, in *Metric) (*empty.Empty, error) {
-	conn, err := c.redis.Connect()
-	if err != nil {
-		return nil, err
-	}
-	defer conn.Close()
-
-	exists, err := redis.Int(conn.Do("SISMEMBER", "FILES", in.Filename))
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to check for file presence")
-	} else if exists == 0 {
-		return nil, status.Error(codes.FailedPrecondition,
-			"file does not exist")
-	}
-
-	exists, err = redis.Int(conn.Do("SADD", "TRACKED_FILES", in.Filename))
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to add file to metrics")
-	} else if exists == 0 {
-		return nil, status.Error(codes.AlreadyExists,
-			"file already is in metrics")
-	}
-
-	return &empty.Empty{}, nil
+	return &empty.Empty{}, c.redis.AddTrackedFile(in.Filename)
 }
 
 func (c *CLI) DelMetric(ctx context.Context, in *Metric) (*empty.Empty, error) {
-	conn, err := c.redis.Connect()
-	if err != nil {
-		return nil, err
-	}
-	defer conn.Close()
-
-	exists, err := redis.Int(conn.Do("SISMEMBER", "TRACKED_FILES", in.Filename))
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to check for file presence")
-	} else if exists == 0 {
-		return nil, status.Error(codes.FailedPrecondition,
-			"file is not part of tracked files")
-	}
-
-	_, err = conn.Do("SREM", "TRACKED_FILES", in.Filename)
-	if err != nil {
-		return nil, errors.Wrap(err, "failed to remove file from tracked files")
-	}
-
-	return &empty.Empty{}, nil
+	return &empty.Empty{}, c.redis.DeleteTrackedFile(in.Filename)
 }
 
 func (c *CLI) ListMetrics(ctx context.Context, in *Metric) (*MetricsList, error) {
@@ -799,3 +758,17 @@ func (c *CLI) ListMetrics(ctx context.Context, in *Metric) (*MetricsList, error)
 
 	return &MetricsList{Filename: files}, nil
 }
+
+func (c *CLI) EnableAuto(context.Context, *empty.Empty) (*empty.Empty, error) {
+	GetConfig().MetricsAutoTrackedFiles = true
+	return &empty.Empty{}, nil
+}
+
+func (c *CLI) DisableAuto(context.Context, *empty.Empty) (*empty.Empty, error) {
+	GetConfig().MetricsAutoTrackedFiles = false
+	return &empty.Empty{}, nil
+}
+
+func (c *CLI) GetStatusAuto(context.Context, *empty.Empty) (*StatusAuto, error) {
+	return &StatusAuto{Status: GetConfig().MetricsAutoTrackedFiles}, nil
+}