From d1dd7c5593fba2f0bdb8592f2420c52bd81db49f Mon Sep 17 00:00:00 2001
From: Fabian Reinartz <fab.reinartz@gmail.com>
Date: Fri, 22 Jan 2016 12:06:16 +0100
Subject: [PATCH] Implement millisecond duration literals

---
 model/time.go      | 49 +++++++++++++++++++++++++++-------------------
 model/time_test.go | 37 ++++++++++++++++++++++++++++++++++
 2 files changed, 66 insertions(+), 20 deletions(-)

diff --git a/model/time.go b/model/time.go
index ebc8bf6cc..e2e9ff574 100644
--- a/model/time.go
+++ b/model/time.go
@@ -163,6 +163,8 @@ func (t *Time) UnmarshalJSON(b []byte) error {
 // This type should not propagate beyond the scope of input/output processing.
 type Duration time.Duration
 
+var durationRE = regexp.MustCompile("^([0-9]+)(d|h|m|s|ms)$")
+
 // StringToDuration parses a string into a time.Duration, assuming that a year
 // a day always has 24h.
 func ParseDuration(durationStr string) (Duration, error) {
@@ -170,44 +172,51 @@ func ParseDuration(durationStr string) (Duration, error) {
 	if len(matches) != 3 {
 		return 0, fmt.Errorf("not a valid duration string: %q", durationStr)
 	}
-	durSeconds, _ := strconv.Atoi(matches[1])
-	dur := time.Duration(durSeconds) * time.Second
-	unit := matches[2]
-	switch unit {
+	var (
+		n, _ = strconv.Atoi(matches[1])
+		dur  = time.Duration(n) * time.Millisecond
+	)
+	switch unit := matches[2]; unit {
 	case "d":
-		dur *= 60 * 60 * 24
+		dur *= 1000 * 60 * 60 * 24
 	case "h":
-		dur *= 60 * 60
+		dur *= 1000 * 60 * 60
 	case "m":
-		dur *= 60
+		dur *= 1000 * 60
 	case "s":
-		dur *= 1
+		dur *= 1000
+	case "ms":
+		// Value already correct
 	default:
 		return 0, fmt.Errorf("invalid time unit in duration string: %q", unit)
 	}
 	return Duration(dur), nil
 }
 
-var durationRE = regexp.MustCompile("^([0-9]+)([ywdhms]+)$")
-
 func (d Duration) String() string {
-	seconds := int64(time.Duration(d) / time.Second)
+	var (
+		ms   = int64(time.Duration(d) / time.Millisecond)
+		unit = "ms"
+	)
 	factors := map[string]int64{
-		"d": 60 * 60 * 24,
-		"h": 60 * 60,
-		"m": 60,
-		"s": 1,
+		"d":  1000 * 60 * 60 * 24,
+		"h":  1000 * 60 * 60,
+		"m":  1000 * 60,
+		"s":  1000,
+		"ms": 1,
 	}
-	unit := "s"
+
 	switch int64(0) {
-	case seconds % factors["d"]:
+	case ms % factors["d"]:
 		unit = "d"
-	case seconds % factors["h"]:
+	case ms % factors["h"]:
 		unit = "h"
-	case seconds % factors["m"]:
+	case ms % factors["m"]:
 		unit = "m"
+	case ms % factors["s"]:
+		unit = "s"
 	}
-	return fmt.Sprintf("%v%v", seconds/factors[unit], unit)
+	return fmt.Sprintf("%v%v", ms/factors[unit], unit)
 }
 
 // MarshalYAML implements the yaml.Marshaler interface.
diff --git a/model/time_test.go b/model/time_test.go
index 9013a6277..2693d1321 100644
--- a/model/time_test.go
+++ b/model/time_test.go
@@ -84,3 +84,40 @@ func TestDuration(t *testing.T) {
 		t.Fatalf("Expected %s to be equal to %s", delta, duration)
 	}
 }
+
+func TestParseDuration(t *testing.T) {
+	var cases = []struct {
+		in  string
+		out time.Duration
+	}{
+		{
+			in:  "324ms",
+			out: 324 * time.Millisecond,
+		}, {
+			in:  "3s",
+			out: 3 * time.Second,
+		}, {
+			in:  "5m",
+			out: 5 * time.Minute,
+		}, {
+			in:  "1h",
+			out: time.Hour,
+		}, {
+			in:  "4d",
+			out: 4 * 24 * time.Hour,
+		},
+	}
+
+	for _, c := range cases {
+		d, err := ParseDuration(c.in)
+		if err != nil {
+			t.Errorf("Unexpected error on input %q", c.in)
+		}
+		if time.Duration(d) != c.out {
+			t.Errorf("Expected %v but got %v", c.out, d)
+		}
+		if d.String() != c.in {
+			t.Errorf("Expected duration string %q but got %q", c.in, d.String())
+		}
+	}
+}