wal-g

1
package ioextensions
2

3
import (
4
	"io"
5

6
	"github.com/wal-g/tracelog"
7
)
8

9
type ReaderWithRetry struct {
10
	reader        io.ReadCloser
11
	getReader     func() (io.ReadCloser, error)
12
	retryAttempts int
13
	alreadyRead   int
14
	attempt       int
15
}
16

17
func NewReaderWithRetry(getReader func() (io.ReadCloser, error), retryAttempts int) io.ReadCloser {
18
	if retryAttempts <= 0 {
19
		retryAttempts = 1
20
	}
21
	return &ReaderWithRetry{
22
		reader:        nil,
23
		getReader:     getReader,
24
		retryAttempts: retryAttempts,
25
		alreadyRead:   0,
26
		attempt:       0,
27
	}
28
}
29

30
func (r *ReaderWithRetry) setupNewReader() error {
31
	reader, err := r.getReader()
32
	if err != nil {
33
		return err
34
	}
35
	_, err = io.CopyN(io.Discard, reader, int64(r.alreadyRead))
36
	r.reader = reader
37
	return err
38
}
39

40
func (r *ReaderWithRetry) Read(p []byte) (int, error) {
41
	n := 0
42
	var lastErr error
43
	for r.attempt < r.retryAttempts {
44
		if r.reader == nil {
45
			err := r.setupNewReader()
46
			if err == io.EOF {
47
				return n, err
48
			} else if err != nil {
49
				tracelog.ErrorLogger.Printf("error while initializing reader: %v", err)
50
				tracelog.ErrorLogger.PrintOnError(r.reader.Close())
51
				r.reader = nil
52
				r.attempt++
53
				continue
54
			}
55
		}
56

57
		read, err := r.reader.Read(p[n:])
58
		lastErr = err
59
		n += read
60
		r.alreadyRead += read
61

62
		if err == io.EOF {
63
			return n, err
64
		} else if err != nil {
65
			tracelog.ErrorLogger.Printf("error while read file: %v. Attempt: %d\n", err, r.attempt)
66
			tracelog.ErrorLogger.PrintOnError(r.reader.Close())
67
			r.reader = nil
68
			r.attempt++
69
			continue
70
		} else if n == len(p) {
71
			return n, nil
72
		}
73
	}
74

75
	return n, lastErr
76
}
77

78
func (r *ReaderWithRetry) Close() error {
79
	if r.reader != nil {
80
		return r.reader.Close()
81
	}
82
	return nil
83
}
84