From 9340be717f34c908ed64c0655679314742813e27 Mon Sep 17 00:00:00 2001
From: Maksym Pavlenko <pavlenko.maksym@gmail.com>
Date: Thu, 1 Feb 2024 11:48:33 -0800
Subject: [PATCH] Remove duplicated TOML duration parsers

Signed-off-by: Maksym Pavlenko <pavlenko.maksym@gmail.com>
---
 plugins/gc/scheduler.go      | 24 +++++-------------------
 plugins/gc/scheduler_test.go |  3 ++-
 plugins/restart/monitor.go   | 23 ++++-------------------
 3 files changed, 11 insertions(+), 39 deletions(-)

diff --git a/plugins/gc/scheduler.go b/plugins/gc/scheduler.go
index a673a25cd..3213ba3dd 100644
--- a/plugins/gc/scheduler.go
+++ b/plugins/gc/scheduler.go
@@ -23,6 +23,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/containerd/containerd/v2/internal/tomlext"
 	"github.com/containerd/containerd/v2/pkg/gc"
 	"github.com/containerd/containerd/v2/plugins"
 	"github.com/containerd/log"
@@ -70,7 +71,7 @@ type config struct {
 	// schedule. Use suffix "ms" for millisecond and "s" for second.
 	//
 	// Default is "0ms"
-	ScheduleDelay duration `toml:"schedule_delay"`
+	ScheduleDelay tomlext.Duration `toml:"schedule_delay"`
 
 	// StartupDelay is the delay duration to do an initial garbage
 	// collection after startup. The initial garbage collection is used to
@@ -79,22 +80,7 @@ type config struct {
 	// "ms" for millisecond and "s" for second.
 	//
 	// Default is "100ms"
-	StartupDelay duration `toml:"startup_delay"`
-}
-
-type duration time.Duration
-
-func (d *duration) UnmarshalText(text []byte) error {
-	ed, err := time.ParseDuration(string(text))
-	if err != nil {
-		return err
-	}
-	*d = duration(ed)
-	return nil
-}
-
-func (d duration) MarshalText() (text []byte, err error) {
-	return []byte(time.Duration(d).String()), nil
+	StartupDelay tomlext.Duration `toml:"startup_delay"`
 }
 
 func init() {
@@ -108,8 +94,8 @@ func init() {
 			PauseThreshold:    0.02,
 			DeletionThreshold: 0,
 			MutationThreshold: 100,
-			ScheduleDelay:     duration(0),
-			StartupDelay:      duration(100 * time.Millisecond),
+			ScheduleDelay:     tomlext.FromStdTime(0),
+			StartupDelay:      tomlext.FromStdTime(100 * time.Millisecond),
 		},
 		InitFn: func(ic *plugin.InitContext) (interface{}, error) {
 			md, err := ic.GetSingle(plugins.MetadataPlugin)
diff --git a/plugins/gc/scheduler_test.go b/plugins/gc/scheduler_test.go
index befd37540..643d37630 100644
--- a/plugins/gc/scheduler_test.go
+++ b/plugins/gc/scheduler_test.go
@@ -22,6 +22,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/containerd/containerd/v2/internal/tomlext"
 	"github.com/containerd/containerd/v2/pkg/gc"
 	"github.com/stretchr/testify/assert"
 )
@@ -152,7 +153,7 @@ func TestStartupDelay(t *testing.T) {
 		cfg          = &config{
 			// Prevent GC from scheduling again before check
 			PauseThreshold: 0.001,
-			StartupDelay:   duration(startupDelay),
+			StartupDelay:   tomlext.Duration(startupDelay),
 		}
 		tc = &testCollector{
 			d: time.Second,
diff --git a/plugins/restart/monitor.go b/plugins/restart/monitor.go
index b47442025..4016159e1 100644
--- a/plugins/restart/monitor.go
+++ b/plugins/restart/monitor.go
@@ -25,6 +25,7 @@ import (
 
 	containerd "github.com/containerd/containerd/v2/client"
 	"github.com/containerd/containerd/v2/core/runtime/restart"
+	"github.com/containerd/containerd/v2/internal/tomlext"
 	"github.com/containerd/containerd/v2/pkg/namespaces"
 	"github.com/containerd/containerd/v2/plugins"
 	"github.com/containerd/log"
@@ -32,24 +33,10 @@ import (
 	"github.com/containerd/plugin/registry"
 )
 
-type duration struct {
-	time.Duration
-}
-
-func (d *duration) UnmarshalText(text []byte) error {
-	var err error
-	d.Duration, err = time.ParseDuration(string(text))
-	return err
-}
-
-func (d duration) MarshalText() ([]byte, error) {
-	return []byte(d.Duration.String()), nil
-}
-
 // Config for the restart monitor
 type Config struct {
 	// Interval for how long to wait to check for state changes
-	Interval duration `toml:"interval"`
+	Interval tomlext.Duration `toml:"interval"`
 }
 
 func init() {
@@ -61,9 +48,7 @@ func init() {
 		},
 		ID: "restart",
 		Config: &Config{
-			Interval: duration{
-				Duration: 10 * time.Second,
-			},
+			Interval: tomlext.FromStdTime(10 * time.Second),
 		},
 		InitFn: func(ic *plugin.InitContext) (interface{}, error) {
 			ic.Meta.Capabilities = []string{"no", "always", "on-failure", "unless-stopped"}
@@ -74,7 +59,7 @@ func init() {
 			m := &monitor{
 				client: client,
 			}
-			go m.run(ic.Config.(*Config).Interval.Duration)
+			go m.run(tomlext.ToStdTime(ic.Config.(*Config).Interval))
 			return m, nil
 		},
 	})