Fiddling with flags to cache server, mostly changing it such that cle…

…an frequency is specified as a 'nice' duration string, e.g. "3h"

src/cache/server/cache.go

@@ -13,6 +13,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/dustin/go-humanize"
 	"github.com/streamrail/concurrent-map"
 
 	"cache/tools"
@@ -41,11 +42,11 @@ type Cache struct {
 // NewCache initialises the cache and fires off a background cleaner goroutine which runs every
 // cleanFrequency seconds. The high and low water marks control a (soft) max size and a (harder)
 // minimum size.
-func NewCache(path string, cleanFrequency int, lowWaterMark, highWaterMark int64) *Cache {
-	log.Notice("Initialising cache with settings:\n  Path: %s\n  Clean frequency: %d\n  Low water mark: %d\n  High water mark: %d",
-		path, cleanFrequency, lowWaterMark, highWaterMark)
+func NewCache(path string, cleanFrequency time.Duration, lowWaterMark, highWaterMark uint64) *Cache {
+	log.Notice("Initialising cache with settings:\n  Path: %s\n  Clean frequency: %s\n  Low water mark: %s\n  High water mark: %s",
+		path, cleanFrequency, humanize.Bytes(lowWaterMark), humanize.Bytes(highWaterMark))
 	cache := newCache(path)
-	go cache.clean(cleanFrequency, lowWaterMark, highWaterMark)
+	go cache.clean(cleanFrequency, int64(lowWaterMark), int64(highWaterMark))
 	return cache
 }
 
@@ -272,8 +273,8 @@ func (cache *Cache) DeleteAllArtifacts() error {
 }
 
 // clean implements a periodic clean of the cache to remove old artifacts.
-func (cache *Cache) clean(cleanFrequency int, lowWaterMark, highWaterMark int64) {
-	for range time.NewTicker(time.Duration(cleanFrequency) * time.Second).C {
+func (cache *Cache) clean(cleanFrequency time.Duration, lowWaterMark, highWaterMark int64) {
+	for range time.NewTicker(cleanFrequency).C {
 		cache.singleClean(lowWaterMark, highWaterMark)
 	}
 }

src/cache/server/http_server_main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"fmt"
 	"net/http"
+	"time"
 
 	"gopkg.in/op/go-logging.v1"
 
@@ -13,20 +14,24 @@ import (
 var log = logging.MustGetLogger("http_cache_server")
 
 var opts struct {
-	Verbosity      int             `short:"v" long:"verbosity" description:"Verbosity of output (higher number = more output, default 2 -> notice, warnings and errors only)" default:"2"`
-	Port           int             `short:"p" long:"port" description:"Port to serve on" default:"8080"`
-	Dir            string          `short:"d" long:"dir" description:"Directory to write into" default:"plz-http-cache"`
-	LowWaterMark   output.ByteSize `short:"l" long:"low_water_mark" description:"Size of cache to clean down to" default:"18G"`
-	HighWaterMark  output.ByteSize `short:"i" long:"high_water_mark" description:"Max size of cache to clean at" default:"20G"`
-	CleanFrequency int             `short:"f" long:"clean_frequency" description:"Frequency to clean cache at, in seconds" default:"10"`
-	LogFile        string          `long:"log_file" description:"File to log to (in addition to stdout)"`
+	Verbosity int    `short:"v" long:"verbosity" description:"Verbosity of output (higher number = more output, default 2 -> notice, warnings and errors only)" default:"2"`
+	Port      int    `short:"p" long:"port" description:"Port to serve on" default:"8080"`
+	Dir       string `short:"d" long:"dir" description:"Directory to write into" default:"plz-http-cache"`
+	LogFile   string `long:"log_file" description:"File to log to (in addition to stdout)"`
+
+	CleanFlags struct {
+		LowWaterMark   output.ByteSize `short:"l" long:"low_water_mark" description:"Size of cache to clean down to" default:"18G"`
+		HighWaterMark  output.ByteSize `short:"i" long:"high_water_mark" description:"Max size of cache to clean at" default:"20G"`
+		CleanFrequency output.Duration `short:"f" long:"clean_frequency" description:"Frequency to clean cache at" default:"10m"`
+	} `group:"Options controlling when to clean the cache"`
 }
 
 func main() {
 	output.ParseFlagsOrDie("Please RPC cache server", &opts)
 	output.InitLogging(opts.Verbosity, opts.LogFile, opts.Verbosity)
 	log.Notice("Initialising cache server...")
-	cache := server.NewCache(opts.Dir, opts.CleanFrequency, int64(opts.LowWaterMark), int64(opts.HighWaterMark))
+	cache := server.NewCache(opts.Dir, time.Duration(opts.CleanFlags.CleanFrequency),
+		uint64(opts.CleanFlags.LowWaterMark), uint64(opts.CleanFlags.HighWaterMark))
 	log.Notice("Starting up http cache server on port %d...", opts.Port)
 	router := server.BuildRouter(cache)
 	http.Handle("/", router)

src/cache/server/rpc_server_main.go

@@ -1,6 +1,8 @@
 package main
 
 import (
+	"time"
+
 	"gopkg.in/op/go-logging.v1"
 
 	"cache/server"
@@ -10,30 +12,38 @@ import (
 var log = logging.MustGetLogger("rpc_cache_server")
 
 var opts struct {
-	Verbosity      int             `short:"v" long:"verbosity" description:"Verbosity of output (higher number = more output, default 2 -> notice, warnings and errors only)" default:"2"`
-	Port           int             `short:"p" long:"port" description:"Port to serve on" default:"7677"`
-	Dir            string          `short:"d" long:"dir" description:"Directory to write into" default:"plz-rpc-cache"`
-	LowWaterMark   output.ByteSize `short:"l" long:"low_water_mark" description:"Size of cache to clean down to" default:"18G"`
-	HighWaterMark  output.ByteSize `short:"i" long:"high_water_mark" description:"Max size of cache to clean at" default:"20G"`
-	CleanFrequency int             `short:"f" long:"clean_frequency" description:"Frequency to clean cache at, in seconds" default:"10"`
-	LogFile        string          `long:"log_file" description:"File to log to (in addition to stdout)"`
-	KeyFile        string          `long:"key_file" description:"File containing PEM-encoded private key."`
-	CertFile       string          `long:"cert_file" description:"File containing PEM-encoded certificate"`
-	CACertFile     string          `long:"ca_cert_file" description:"File containing PEM-encoded CA certificate"`
-	WritableCerts  string          `long:"writable_certs" description:"File or directory containing certificates that are allowed to write to the cache"`
-	ReadonlyCerts  string          `long:"readonly_certs" description:"File or directory containing certificates that are allowed to read from the cache"`
+	Port      int    `short:"p" long:"port" description:"Port to serve on" default:"7677"`
+	Dir       string `short:"d" long:"dir" description:"Directory to write into" default:"plz-rpc-cache"`
+	Verbosity int    `short:"v" long:"verbosity" description:"Verbosity of output (higher number = more output, default 2 -> notice, warnings and errors only)" default:"2"`
+	LogFile   string `long:"log_file" description:"File to log to (in addition to stdout)"`
+
+	CleanFlags struct {
+		LowWaterMark   output.ByteSize `short:"l" long:"low_water_mark" description:"Size of cache to clean down to" default:"18G"`
+		HighWaterMark  output.ByteSize `short:"i" long:"high_water_mark" description:"Max size of cache to clean at" default:"20G"`
+		CleanFrequency output.Duration `short:"f" long:"clean_frequency" description:"Frequency to clean cache at" default:"10m"`
+	} `group:"Options controlling when to clean the cache"`
+
+	TLSFlags struct {
+		KeyFile       string `long:"key_file" description:"File containing PEM-encoded private key."`
+		CertFile      string `long:"cert_file" description:"File containing PEM-encoded certificate"`
+		CACertFile    string `long:"ca_cert_file" description:"File containing PEM-encoded CA certificate"`
+		WritableCerts string `long:"writable_certs" description:"File or directory containing certificates that are allowed to write to the cache"`
+		ReadonlyCerts string `long:"readonly_certs" description:"File or directory containing certificates that are allowed to read from the cache"`
+	} `group:"Options controlling TLS communication & authentication"`
 }
 
 func main() {
 	output.ParseFlagsOrDie("Please RPC cache server", &opts)
 	output.InitLogging(opts.Verbosity, opts.LogFile, opts.Verbosity)
-	if (opts.KeyFile == "") != (opts.CertFile == "") {
+	if (opts.TLSFlags.KeyFile == "") != (opts.TLSFlags.CertFile == "") {
 		log.Fatalf("Must pass both --key_file and --cert_file if you pass one")
-	} else if opts.KeyFile == "" && (opts.WritableCerts != "" || opts.ReadonlyCerts != "") {
+	} else if opts.TLSFlags.KeyFile == "" && (opts.TLSFlags.WritableCerts != "" || opts.TLSFlags.ReadonlyCerts != "") {
 		log.Fatalf("You can only use --writable_certs / --readonly_certs with https (--key_file and --cert_file)")
 	}
 	log.Notice("Scanning existing cache directory %s...", opts.Dir)
-	cache := server.NewCache(opts.Dir, opts.CleanFrequency, int64(opts.LowWaterMark), int64(opts.HighWaterMark))
+	cache := server.NewCache(opts.Dir, time.Duration(opts.CleanFlags.CleanFrequency),
+		uint64(opts.CleanFlags.LowWaterMark), uint64(opts.CleanFlags.HighWaterMark))
 	log.Notice("Starting up RPC cache server on port %d...", opts.Port)
-	server.ServeGrpcForever(opts.Port, cache, opts.KeyFile, opts.CertFile, opts.CACertFile, opts.ReadonlyCerts, opts.WritableCerts)
+	server.ServeGrpcForever(opts.Port, cache, opts.TLSFlags.KeyFile, opts.TLSFlags.CertFile,
+		opts.TLSFlags.CACertFile, opts.TLSFlags.ReadonlyCerts, opts.TLSFlags.WritableCerts)
 }

src/output/BUILD

@@ -70,3 +70,12 @@ go_test(
         '//src/core',
     ],
 )
+
+go_test(
+    name = 'flags_test',
+    srcs = ['flags_test.go'],
+    deps = [
+        ':output',
+        '//third_party/go:testify',
+    ],
+)

src/output/flags.go

@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"os"
 	"path"
+	"strconv"
 	"strings"
+	"time"
 
 	"github.com/dustin/go-humanize"
 	"github.com/jessevdk/go-flags"
@@ -61,10 +63,29 @@ func ParseFlagsFromArgsOrDie(appname string, data interface{}, args []string) *f
 
 // A ByteSize is used for flags that represent some quantity of bytes that can be
 // passed as human-readable quantities (eg. "10G").
-type ByteSize int64
+type ByteSize uint64
 
+// UnmarshalFlag implements the flags.Unmarshaler interface.
 func (b *ByteSize) UnmarshalFlag(in string) error {
 	b2, err := humanize.ParseBytes(in)
 	*b = ByteSize(b2)
 	return err
 }
+
+// A Duration is used for flags that represent a time duration; it's just a wrapper
+// around time.Duration that implements the flags.Unmarshaler interface.
+type Duration time.Duration
+
+// UnmarshalFlag implements the flags.Unmarshaler interface.
+func (d *Duration) UnmarshalFlag(in string) error {
+	d2, err := time.ParseDuration(in)
+	// For backwards compatibility, treat missing units as seconds.
+	if err != nil {
+		if d3, err := strconv.Atoi(in); err == nil {
+			*d = Duration(time.Duration(d3) * time.Second)
+			return nil
+		}
+	}
+	*d = Duration(d2)
+	return err
+}

src/output/flags_test.go

@@ -0,0 +1,43 @@
+package output
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestByteSize(t *testing.T) {
+	opts := struct {
+		Size ByteSize `short:"b"`
+	}{}
+	_, extraArgs, err := ParseFlags("test", &opts, []string{"test", "-b=15M"})
+	assert.NoError(t, err)
+	assert.Equal(t, 0, len(extraArgs))
+	assert.EqualValues(t, 15000000, opts.Size)
+}
+
+func TestDuration(t *testing.T) {
+	opts := struct {
+		D Duration `short:"d"`
+	}{}
+	_, extraArgs, err := ParseFlags("test", &opts, []string{"test", "-d=3h"})
+	assert.NoError(t, err)
+	assert.Equal(t, 0, len(extraArgs))
+	assert.EqualValues(t, 3*time.Hour, opts.D)
+
+	_, extraArgs, err = ParseFlags("test", &opts, []string{"test", "-d=3"})
+	assert.NoError(t, err)
+	assert.Equal(t, 0, len(extraArgs))
+	assert.EqualValues(t, 3*time.Second, opts.D)
+}
+
+func TestDurationDefault(t *testing.T) {
+	opts := struct {
+		D Duration `short:"d" default:"3h"`
+	}{}
+	_, extraArgs, err := ParseFlags("test", &opts, []string{"test"})
+	assert.NoError(t, err)
+	assert.Equal(t, 0, len(extraArgs))
+	assert.EqualValues(t, 3*time.Hour, opts.D)
+}