Metrics: Add auto tracked files mode

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

Author: Skantes

cli/commands.go

@@ -101,9 +101,11 @@ func (c *cli) CmdHelp() error {
 		{"add", "Add a new mirror"},
 		{"addMetric", "Add a tracked file to the metrics route"},
 		{"disable", "Disable a mirror"},
+		{"disableAuto", "Disable automatic file tracking"},
 		{"delMetric", "Delete a tracked file from the metrics route"},
 		{"edit", "Edit a mirror"},
 		{"enable", "Enable a mirror"},
+		{"enableAuto", "Enable automatic file tracking"},
 		{"export", "Export the mirror database"},
 		{"list", "List all mirrors"},
 		{"listMetrics", "List all tracked files from the metrics route"},
@@ -114,6 +116,7 @@ func (c *cli) CmdHelp() error {
 		{"scan", "(Re-)Scan a mirror"},
 		{"show", "Print a mirror configuration"},
 		{"stats", "Show download stats"},
+		{"statusAuto", "Show automatic file tracking boolean"},
 		{"upgrade", "Seamless binary upgrade"},
 		{"version", "Print version information"},
 	} {
@@ -444,6 +447,69 @@ func (c *cli) CmdListmetrics(args ...string) error {
 	return nil
 }
 
+func (c *cli) CmdEnableauto(args ...string) error {
+	cmd := SubCmd("enableAuto", "", "Enable automatic addition of new files to tracked files")
+
+	if err := cmd.Parse(args); err != nil {
+		return nil
+	}
+	if cmd.NArg() != 0 {
+		cmd.Usage()
+		return nil
+	}
+
+	client := c.GetRPC()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultRPCTimeout)
+	defer cancel()
+	client.EnableAuto(ctx, &empty.Empty{})
+
+	return nil
+}
+
+func (c *cli) CmdDisableauto(args ...string) error {
+	cmd := SubCmd("disableAuto", "", "Disable automatic addition of new files to tracked files")
+
+	if err := cmd.Parse(args); err != nil {
+		return nil
+	}
+	if cmd.NArg() != 0 {
+		cmd.Usage()
+		return nil
+	}
+
+	client := c.GetRPC()
+	ctx, cancel := context.WithTimeout(context.Background(), defaultRPCTimeout)
+	defer cancel()
+	client.DisableAuto(ctx, &empty.Empty{})
+
+	return nil
+}
+
+func (c *cli) CmdStatusauto(args ...string) error {
+	cmd := SubCmd("statusAuto", "", "Print boolean of automatic addition of new files to tracked files")
+
+	if err := cmd.Parse(args); err != nil {
+		return nil
+	}
+	if cmd.NArg() != 0 {
+		cmd.Usage()
+		return nil
+	}
+
+	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

@@ -60,6 +60,7 @@ func defaultConfig() Configuration {
 		RPCListenAddress:         "localhost:3390",
 		RPCPassword:              "",
 		MetricsTopFilesRetention: 1,
+		MetricsAutoTrackedFiles:  false,
 	}
 }
 
@@ -95,7 +96,8 @@ type Configuration struct {
 	RPCListenAddress string `yaml:"RPCListenAddress"`
 	RPCPassword      string `yaml:"RPCPassword"`
 
-	MetricsTopFilesRetention int `yaml:"MetricsTopFilesRetention"`
+	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) {
@@ -160,6 +162,52 @@ func (r *Redis) AddCountry(country string) error {
 	return nil
 }
 
+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

@@ -137,3 +137,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: 1
+
+## 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
+}