Makes wayback timeout configurable

waybackarchiver · waybackarchiver · commit b6f3e8a9a395 · 2021-10-17T23:33:58.000+08:00
diff --git a/README.md b/README.md
@@ -210,6 +210,7 @@ You can also specify configuration options either via command flags or via envir
 | -                   | `WAYBACK_BOLT_PATH`               | `./wayback.db`          | File path of bolt database                                   |
 | -                   | `WAYBACK_STORAGE_DIR`             | -                       | Directory to store binary file, e.g. PDF, html file          |
 | -                   | `WAYBACK_MAX_MEDIA_SIZE`          | `512MB`                 | Max size to limit download stream media                      |
+| -                   | `WAYBACK_TIMEOUT`                 | `300`                   | Timeout for single wayback request, defaults to 300 second   |
 | `-d`, `--daemon`    | -                                 | -                       | Run as daemon service, e.g. `telegram`, `web`, `mastodon`, `twitter`, `discord` |
 | `--ia`              | `WAYBACK_ENABLE_IA`               | `true`                  | Wayback webpages to **Internet Archive**                     |
 | `--is`              | `WAYBACK_ENABLE_IS`               | `true`                  | Wayback webpages to **Archive Today**                        |
diff --git a/config/config_test.go b/config/config_test.go
@@ -12,6 +12,7 @@ import (
 	"os"
 	"strconv"
 	"testing"
+	"time"
 
 	"github.com/wabarc/logger"
 )
@@ -1210,3 +1211,39 @@ func TestMaxMediaSize(t *testing.T) {
 		})
 	}
 }
+
+func TestWaybackTimeout(t *testing.T) {
+	t.Parallel()
+
+	var tests = []struct {
+		timeout  int
+		expected time.Duration
+	}{
+		{
+			timeout:  0,
+			expected: 0 * time.Second,
+		},
+		{
+			timeout:  1,
+			expected: time.Second,
+		},
+	}
+
+	for i, test := range tests {
+		t.Run(strconv.Itoa(i), func(t *testing.T) {
+			os.Clearenv()
+			os.Setenv("WAYBACK_TIMEOUT", strconv.Itoa(test.timeout))
+
+			parser := NewParser()
+			opts, err := parser.ParseEnvironmentVariables()
+			if err != nil {
+				t.Fatalf(`Parsing environment variables failed: %v`, err)
+			}
+
+			got := opts.WaybackTimeout()
+			if got != test.expected {
+				t.Fatalf(`Unexpected set wayback timeout got %d instead of %d`, got, test.expected)
+			}
+		})
+	}
+}
diff --git a/config/options.go b/config/options.go
@@ -7,6 +7,7 @@ package config // import "github.com/wabarc/wayback/config"
 import (
 	"net/url"
 	"strings"
+	"time"
 
 	"github.com/dustin/go-humanize"
 	"github.com/wabarc/logger"
@@ -71,6 +72,7 @@ const (
 	defPoolingSize         = 3
 	defStorageDir          = ""
 	defMaxMediaSize        = "512MB"
+	defWaybackTimeout      = 300
 )
 
 var (
@@ -104,6 +106,7 @@ type Options struct {
 	poolingSize         int
 	storageDir          string
 	maxMediaSize        string
+	waybackTimeout      int
 }
 
 type ipfs struct {
@@ -189,6 +192,7 @@ func NewOptions() *Options {
 		poolingSize:         defPoolingSize,
 		storageDir:          defStorageDir,
 		maxMediaSize:        defMaxMediaSize,
+		waybackTimeout:      defWaybackTimeout,
 		ipfs: &ipfs{
 			host: defIPFSHost,
 			port: defIPFSPort,
@@ -621,3 +625,8 @@ func (o *Options) MaxAttachSize(scope string) int64 {
 	}
 	return scopes[scope]
 }
+
+// WaybackTimeout returns timeout for a wayback request.
+func (o *Options) WaybackTimeout() time.Duration {
+	return time.Duration(o.waybackTimeout) * time.Second
+}
diff --git a/config/parser.go b/config/parser.go
@@ -183,6 +183,8 @@ func (p *Parser) parseLines(lines []string) (err error) {
 			p.opts.storageDir = parseString(val, defStorageDir)
 		case "WAYBACK_MAX_MEDIA_SIZE":
 			p.opts.maxMediaSize = parseString(val, defMaxMediaSize)
+		case "WAYBACK_TIMEOUT":
+			p.opts.waybackTimeout = parseInt(val, defWaybackTimeout)
 		default:
 			if os.Getenv(key) == "" && val != "" {
 				os.Setenv(key, val)
diff --git a/wayback.1 b/wayback.1
@@ -153,6 +153,9 @@ Enable IPFS\&.
 .B WAYBACK_POOLING_SIZE
 Number of worker pool for wayback at once. default 3\&.
 .TP
+.B WAYBACK_TIMEOUT
+Timeout for single wayback request, default 300\&.
+.TP
 .B WAYBACK_BOLT_PATH
 File path of bolt database. default ./wayback.db\&.
 .TP
diff --git a/wayback.conf b/wayback.conf
@@ -48,6 +48,7 @@ CHROME_REMOTE_ADDR=
 WAYBACK_POOLING_SIZE=3
 WAYBACK_STORAGE_DIR=
 WAYBACK_MAX_MEDIA_SIZE=512MB
+WAYBACK_TIMEOUT=300
 
 # ipfs slot: infura, pinata
 # doc: https://github.com/wabarc/ipfs-pinner#supported-pinning-services
diff --git a/wayback.go b/wayback.go
@@ -9,7 +9,6 @@ import (
 	"fmt"
 	"net/url"
 	"sync"
-	"time"
 
 	"github.com/wabarc/logger"
 	"github.com/wabarc/playback"
@@ -150,14 +149,14 @@ func wayback(w Waybacker) string {
 func Wayback(ctx context.Context, bundles *reduxer.Bundles, urls ...string) (cols []Collect, err error) {
 	logger.Debug("start...")
 
+	ctx, cancel := context.WithTimeout(ctx, config.Opts.WaybackTimeout())
+	defer cancel()
+
 	*bundles, err = reduxer.Do(ctx, urls...)
 	if err != nil {
 		logger.Warn("cannot to start reduxer: %v", err)
 	}
 
-	ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
-	defer cancel()
-
 	mu := sync.Mutex{}
 	g, ctx := errgroup.WithContext(ctx)
 	for _, uri := range urls {
@@ -214,6 +213,9 @@ func Wayback(ctx context.Context, bundles *reduxer.Bundles, urls ...string) (col
 func Playback(ctx context.Context, urls ...string) (cols []Collect, err error) {
 	logger.Debug("start...")
 
+	ctx, cancel := context.WithTimeout(ctx, config.Opts.WaybackTimeout())
+	defer cancel()
+
 	mu := sync.Mutex{}
 	g, ctx := errgroup.WithContext(ctx)
 	var slots = []string{config.SLOT_IA, config.SLOT_IS, config.SLOT_IP, config.SLOT_PH, config.SLOT_TT, config.SLOT_GC}
