feat: implement SyncTimeout (#33)

* feat: implement SyncTimeOut

* Wrap s.context() in new context

* feat: use type time.Duration for SyncTimeout

* test: add testing for config parsing

* test: add test for waitForSync()

* test: modify testcase to use background context with cancel

* chore: set sync_timeout example value in dist/daemon.toml

* chore: fmt: Use %s for time.Duration

* chore: Use ctx from param

---------

Co-authored-by: iBug <git@ibugone.com>

Author: taoky

cmd/yukid/README.md

@@ -67,6 +67,11 @@ repo_config_dir = ["/path/to/config-dir"]
 ## 设置更新用到的 docker images 的频率
 ## 格式为 crontab
 #images_upgrade_interval = "@every 1h"
+
+## 同步超时时间，如果超过了这个时间，同步容器会被强制停止
+## 支持使用 time.ParseDuration() 支持的时间格式，诸如 "10m", "1h" 等
+## 如果为 0 的话则不会超时。注意修改的配置仅对新启动的同步容器生效
+#sync_timeout = "48h"
 ```
 
 ### Repo Configuration

dist/daemon.toml

@@ -32,3 +32,5 @@ repo_config_dir = "/path/to/config-dir"
 #post_sync = ["/path/to/the/program"]
 
 #images_upgrade_interval = "@every 1h"
+
+#sync_timeout = "48h"

pkg/server/config.go

@@ -3,6 +3,7 @@ package server
 import (
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/viper"
@@ -19,15 +20,16 @@ type AppConfig struct {
 	FileSystem     string `mapstructure:"fs,omitempty" validate:"omitempty,eq=xfs|eq=zfs|eq=default"`
 	DockerEndpoint string `mapstructure:"docker_endpoint,omitempty" validate:"omitempty,unix_addr|tcp_addr"`
 
-	Owner                 string   `mapstructure:"owner,omitempty" validate:"-"`
-	LogDir                string   `mapstructure:"log_dir,omitempty" validate:"-"`
-	RepoConfigDir         []string   `mapstructure:"repo_config_dir,omitempty" validate:"required"`
-	LogLevel              string   `mapstructure:"log_level,omitempty" validate:"omitempty,eq=debug|eq=info|eq=warn|eq=error"`
-	ListenAddr            string   `mapstructure:"listen_addr,omitempty" validate:"omitempty,hostport"`
-	BindIP                string   `mapstructure:"bind_ip,omitempty" validate:"omitempty,ip"`
-	NamePrefix            string   `mapstructure:"name_prefix,omitempty" validate:"-"`
-	PostSync              []string `mapstructure:"post_sync,omitempty" validate:"-"`
-	ImagesUpgradeInterval string   `mapstructure:"images_upgrade_interval,omitempty" validate:"omitempty,cron"`
+	Owner                 string        `mapstructure:"owner,omitempty" validate:"-"`
+	LogDir                string        `mapstructure:"log_dir,omitempty" validate:"-"`
+	RepoConfigDir         []string      `mapstructure:"repo_config_dir,omitempty" validate:"required"`
+	LogLevel              string        `mapstructure:"log_level,omitempty" validate:"omitempty,eq=debug|eq=info|eq=warn|eq=error"`
+	ListenAddr            string        `mapstructure:"listen_addr,omitempty" validate:"omitempty,hostport"`
+	BindIP                string        `mapstructure:"bind_ip,omitempty" validate:"omitempty,ip"`
+	NamePrefix            string        `mapstructure:"name_prefix,omitempty" validate:"-"`
+	PostSync              []string      `mapstructure:"post_sync,omitempty" validate:"-"`
+	ImagesUpgradeInterval string        `mapstructure:"images_upgrade_interval,omitempty" validate:"omitempty,cron"`
+	SyncTimeout           time.Duration `mapstructure:"sync_timeout,omitempty" validate:"omitempty,gte=0"`
 }
 
 type Config struct {
@@ -41,6 +43,7 @@ type Config struct {
 	NamePrefix            string
 	PostSync              []string
 	ImagesUpgradeInterval string
+	SyncTimeout           time.Duration
 }
 
 var (
@@ -57,6 +60,7 @@ var (
 		NamePrefix:            "syncing-",
 		LogLevel:              logrus.InfoLevel,
 		ImagesUpgradeInterval: "@every 1h",
+		SyncTimeout:           0,
 	}
 )
 
@@ -87,6 +91,7 @@ func LoadConfig() (*Config, error) {
 		NamePrefix:            appCfg.NamePrefix,
 		PostSync:              appCfg.PostSync,
 		ImagesUpgradeInterval: appCfg.ImagesUpgradeInterval,
+		SyncTimeout:           appCfg.SyncTimeout,
 	}
 
 	switch appCfg.FileSystem {

pkg/server/config_test.go

@@ -0,0 +1,29 @@
+package server
+
+import (
+	"testing"
+	"time"
+
+	"github.com/spf13/viper"
+)
+
+func TestDefaultTimeoutConfig(t *testing.T) {
+	viper.Reset()
+	viper.SetConfigFile("../../dist/daemon.toml")
+	_, err := LoadConfig()
+	if err != nil {
+		t.Fatal(err)
+	}
+}
+
+func TestSyncTimeoutConfig(t *testing.T) {
+	viper.Reset()
+	viper.SetConfigFile("../../test/fixtures/sync_timeout.toml")
+	config, err := LoadConfig()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if config.SyncTimeout != time.Second*15 {
+		t.Fatalf("Expected SyncTimeout to be 15s, got %s", config.SyncTimeout)
+	}
+}

pkg/server/main.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -396,8 +397,21 @@ func (s *Server) waitForSync(ct *api.Container) error {
 		Name: ct.Labels["org.ustcmirror.name"],
 	}
 
-	code, err := s.c.WaitContainer(s.context(), ct.ID)
+	ctx := s.context()
+	if s.config.SyncTimeout > 0 {
+		var cancel context.CancelFunc
+		ctx, cancel = context.WithTimeout(ctx, s.config.SyncTimeout)
+		defer cancel()
+	}
+
+	code, err := s.c.WaitContainer(ctx, ct.ID)
 	if err != nil {
+		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+			// timeout, we should stop the container
+			ctx, cancel := context.WithTimeout(s.context(), time.Second*30)
+			_ = s.c.StopContainer(ctx, ct.ID)
+			cancel()
+		}
 		return err
 	}
 

pkg/server/main_test.go

@@ -0,0 +1,78 @@
+package server
+
+import (
+	"context"
+	"os"
+	"testing"
+
+	"github.com/spf13/viper"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"github.com/ustclug/Yuki/pkg/api"
+	"github.com/ustclug/Yuki/pkg/core"
+)
+
+func getTestingServer(ctx context.Context, prefix string, name string, storageDir string) (*Server, error) {
+	s, err := New()
+	if err != nil {
+		return nil, err
+	}
+
+	defer os.RemoveAll("/tmp/log_yuki")
+	defer os.RemoveAll("/tmp/config_yuki")
+
+	err = os.MkdirAll(storageDir, os.ModePerm)
+	if err != nil {
+		return nil, err
+	}
+	defer os.RemoveAll(storageDir)
+	// remove existing repository and container
+	_ = s.c.RemoveRepository(name)
+	_ = s.c.RemoveContainer(ctx, prefix+name)
+
+	s.ctx = ctx
+	go s.onPreSync()
+	go s.onPostSync()
+
+	return s, nil
+}
+
+func TestWaitForSync(t *testing.T) {
+	as := assert.New(t)
+	req := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	viper.Reset()
+	viper.SetConfigFile("../../test/fixtures/sync_timeout.toml")
+
+	prefix := "syncing-"
+	name := "yuki-test-sync-repo"
+	d := "/tmp/" + name
+	cycle := 10
+
+	s, err := getTestingServer(ctx, prefix, name, d)
+	req.Nil(err)
+
+	err = s.c.AddRepository(api.Repository{
+		Name:        name,
+		Interval:    "1 * * * *",
+		Image:       "ustcmirror/test:latest",
+		StorageDir:  d,
+		LogRotCycle: &cycle,
+		Envs:        map[string]string{"SLEEP_INFINITY": "1"},
+	})
+	as.Nil(err)
+
+	ct, err := s.c.Sync(ctx, core.SyncOptions{
+		MountDir:   false,
+		Name:       name,
+		NamePrefix: prefix,
+	})
+	as.Nil(err)
+	defer s.c.RemoveContainer(ctx, ct.ID)
+	err = s.waitForSync(ct)
+	if err != context.DeadlineExceeded {
+		t.Fatalf("Expected error to be context.DeadlineExceeded, got %v", err)
+	}
+}

test/fixtures/normal.toml

@@ -0,0 +1,7 @@
+fs = "default"
+
+db_name = "test"
+
+log_dir = "/tmp/log_yuki/"
+
+repo_config_dir = "/tmp/config_yuki"

test/fixtures/sync_timeout.toml

@@ -0,0 +1,9 @@
+fs = "default"
+
+db_name = "test"
+
+log_dir = "/tmp/log_yuki/"
+
+repo_config_dir = "/tmp/config_yuki"
+
+sync_timeout = "15s"