From 9aa9e0d1ec6c8453406337756f211eaaeaffcf83 Mon Sep 17 00:00:00 2001
From: Michael Eischer <michael.eischer@fau.de>
Date: Thu, 8 Jun 2023 13:06:25 +0200
Subject: [PATCH] local/sftp: move limiter setup into backend

---
 cmd/restic/global.go                        |  5 -----
 internal/backend/limiter/limiter_backend.go | 15 +++++++++++++++
 internal/backend/local/local.go             |  6 +-----
 internal/backend/sftp/sftp.go               |  6 +-----
 4 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/cmd/restic/global.go b/cmd/restic/global.go
index f4886ec90..f08a75b13 100644
--- a/cmd/restic/global.go
+++ b/cmd/restic/global.go
@@ -610,11 +610,6 @@ func open(ctx context.Context, s string, gopts GlobalOptions, opts options.Optio
 		}
 	}
 
-	if loc.Scheme == "local" || loc.Scheme == "sftp" {
-		// wrap the backend in a LimitBackend so that the throughput is limited
-		be = limiter.LimitBackend(be, lim)
-	}
-
 	// check if config is there
 	fi, err := be.Stat(ctx, restic.Handle{Type: restic.ConfigFile})
 	if err != nil {
diff --git a/internal/backend/limiter/limiter_backend.go b/internal/backend/limiter/limiter_backend.go
index 7fcca59cc..a91794037 100644
--- a/internal/backend/limiter/limiter_backend.go
+++ b/internal/backend/limiter/limiter_backend.go
@@ -7,6 +7,21 @@ import (
 	"github.com/restic/restic/internal/restic"
 )
 
+func WrapBackendConstructor[B restic.Backend, C any](constructor func(ctx context.Context, cfg C) (B, error)) func(ctx context.Context, cfg C, lim Limiter) (restic.Backend, error) {
+	return func(ctx context.Context, cfg C, lim Limiter) (restic.Backend, error) {
+		var be restic.Backend
+		be, err := constructor(ctx, cfg)
+		if err != nil {
+			return nil, err
+		}
+
+		if lim != nil {
+			be = LimitBackend(be, lim)
+		}
+		return be, nil
+	}
+}
+
 // LimitBackend wraps a Backend and applies rate limiting to Load() and Save()
 // calls on the backend.
 func LimitBackend(be restic.Backend, l Limiter) restic.Backend {
diff --git a/internal/backend/local/local.go b/internal/backend/local/local.go
index 02ac81b8d..e9d00abf7 100644
--- a/internal/backend/local/local.go
+++ b/internal/backend/local/local.go
@@ -31,11 +31,7 @@ type Local struct {
 var _ restic.Backend = &Local{}
 
 func NewFactory() location.Factory {
-	return location.NewLimitedBackendFactory(ParseConfig, location.NoPassword, func(ctx context.Context, cfg Config, _ limiter.Limiter) (*Local, error) {
-		return Create(ctx, cfg)
-	}, func(ctx context.Context, cfg Config, _ limiter.Limiter) (*Local, error) {
-		return Open(ctx, cfg)
-	})
+	return location.NewLimitedBackendFactory(ParseConfig, location.NoPassword, limiter.WrapBackendConstructor(Create), limiter.WrapBackendConstructor(Open))
 }
 
 const defaultLayout = "default"
diff --git a/internal/backend/sftp/sftp.go b/internal/backend/sftp/sftp.go
index f0a7ef9bc..1e12df808 100644
--- a/internal/backend/sftp/sftp.go
+++ b/internal/backend/sftp/sftp.go
@@ -44,11 +44,7 @@ type SFTP struct {
 var _ restic.Backend = &SFTP{}
 
 func NewFactory() location.Factory {
-	return location.NewLimitedBackendFactory(ParseConfig, location.NoPassword, func(ctx context.Context, cfg Config, _ limiter.Limiter) (*SFTP, error) {
-		return Create(ctx, cfg)
-	}, func(ctx context.Context, cfg Config, _ limiter.Limiter) (*SFTP, error) {
-		return Open(ctx, cfg)
-	})
+	return location.NewLimitedBackendFactory(ParseConfig, location.NoPassword, limiter.WrapBackendConstructor(Create), limiter.WrapBackendConstructor(Open))
 }
 
 const defaultLayout = "default"