From d139faa40c7c6a0e6ed8770e7a721807d9db6476 Mon Sep 17 00:00:00 2001
From: Jason Song <i@wolfogre.com>
Date: Thu, 6 Apr 2023 10:57:36 +0800
Subject: [PATCH] Supports configuring fetch_timeout and fetch_interval. (#100)

Fix #99.

Fix #94.

Reviewed-on: https://gitea.com/gitea/act_runner/pulls/100
Reviewed-by: Zettat123 <zettat123@noreply.gitea.io>
---
 internal/app/poll/poller.go             | 19 +++++++++----------
 internal/pkg/config/config.example.yaml |  4 ++++
 internal/pkg/config/config.go           | 20 ++++++++++++++------
 3 files changed, 27 insertions(+), 16 deletions(-)

diff --git a/internal/app/poll/poller.go b/internal/app/poll/poller.go
index 29c4e33..c201eff 100644
--- a/internal/app/poll/poller.go
+++ b/internal/app/poll/poller.go
@@ -7,7 +7,6 @@ import (
 	"context"
 	"errors"
 	"sync"
-	"time"
 
 	runnerv1 "code.gitea.io/actions-proto-go/runner/v1"
 	"github.com/bufbuild/connect-go"
@@ -20,23 +19,23 @@ import (
 )
 
 type Poller struct {
-	client   client.Client
-	runner   *run.Runner
-	capacity int
+	client client.Client
+	runner *run.Runner
+	cfg    *config.Config
 }
 
 func New(cfg *config.Config, client client.Client, runner *run.Runner) *Poller {
 	return &Poller{
-		client:   client,
-		runner:   runner,
-		capacity: cfg.Runner.Capacity,
+		client: client,
+		runner: runner,
+		cfg:    cfg,
 	}
 }
 
 func (p *Poller) Poll(ctx context.Context) {
-	limiter := rate.NewLimiter(rate.Every(2*time.Second), 1)
+	limiter := rate.NewLimiter(rate.Every(p.cfg.Runner.FetchInterval), 1)
 	wg := &sync.WaitGroup{}
-	for i := 0; i < p.capacity; i++ {
+	for i := 0; i < p.cfg.Runner.Capacity; i++ {
 		wg.Add(1)
 		go p.poll(ctx, wg, limiter)
 	}
@@ -63,7 +62,7 @@ func (p *Poller) poll(ctx context.Context, wg *sync.WaitGroup, limiter *rate.Lim
 }
 
 func (p *Poller) fetchTask(ctx context.Context) (*runnerv1.Task, bool) {
-	reqCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	reqCtx, cancel := context.WithTimeout(ctx, p.cfg.Runner.FetchTimeout)
 	defer cancel()
 
 	resp, err := p.client.FetchTask(reqCtx, connect.NewRequest(&runnerv1.FetchTaskRequest{}))
diff --git a/internal/pkg/config/config.example.yaml b/internal/pkg/config/config.example.yaml
index 871b346..38ffaff 100644
--- a/internal/pkg/config/config.example.yaml
+++ b/internal/pkg/config/config.example.yaml
@@ -22,6 +22,10 @@ runner:
   timeout: 3h
   # Whether skip verifying the TLS certificate of the Gitea instance.
   insecure: false
+  # The timeout for fetching the job from the Gitea instance.
+  fetch_timeout: 5s
+  # The interval for fetching the job from the Gitea instance.
+  fetch_interval: 2s
 
 cache:
   # Enable cache server to use actions/cache.
diff --git a/internal/pkg/config/config.go b/internal/pkg/config/config.go
index 6280f47..4c3380a 100644
--- a/internal/pkg/config/config.go
+++ b/internal/pkg/config/config.go
@@ -18,12 +18,14 @@ type Config struct {
 		Level string `yaml:"level"`
 	} `yaml:"log"`
 	Runner struct {
-		File     string            `yaml:"file"`
-		Capacity int               `yaml:"capacity"`
-		Envs     map[string]string `yaml:"envs"`
-		EnvFile  string            `yaml:"env_file"`
-		Timeout  time.Duration     `yaml:"timeout"`
-		Insecure bool              `yaml:"insecure"`
+		File          string            `yaml:"file"`
+		Capacity      int               `yaml:"capacity"`
+		Envs          map[string]string `yaml:"envs"`
+		EnvFile       string            `yaml:"env_file"`
+		Timeout       time.Duration     `yaml:"timeout"`
+		Insecure      bool              `yaml:"insecure"`
+		FetchTimeout  time.Duration     `yaml:"fetch_timeout"`
+		FetchInterval time.Duration     `yaml:"fetch_interval"`
 	} `yaml:"runner"`
 	Cache struct {
 		Enabled *bool  `yaml:"enabled"` // pointer to distinguish between false and not set, and it will be true if not set
@@ -90,6 +92,12 @@ func LoadDefault(file string) (*Config, error) {
 	if cfg.Container.NetworkMode == "" {
 		cfg.Container.NetworkMode = "bridge"
 	}
+	if cfg.Runner.FetchTimeout <= 0 {
+		cfg.Runner.FetchTimeout = 5 * time.Second
+	}
+	if cfg.Runner.FetchInterval <= 0 {
+		cfg.Runner.FetchInterval = 2 * time.Second
+	}
 
 	return cfg, nil
 }