wal-g

1
package internal
2

3
import (
4
	"archive/tar"
5
	"context"
6
	"fmt"
7
	"io"
8
	"sync/atomic"
9

10
	"github.com/pkg/errors"
11
	"github.com/wal-g/tracelog"
12
	"github.com/wal-g/wal-g/internal/crypto"
13
	"github.com/wal-g/wal-g/utility"
14
)
15

16
const TarPartitionFolderName = "/tar_partitions/"
17

18
// StorageTarBall represents a tar file that is
19
// going to be uploaded to storage.
20
type StorageTarBall struct {
21
	backupName  string
22
	partNumber  int
23
	partSize    *int64
24
	writeCloser io.Closer
25
	tarWriter   *tar.Writer
26
	uploader    Uploader
27
	name        string
28
}
29

30
func (tarBall *StorageTarBall) Name() string {
31
	return tarBall.name
32
}
33

34
// SetUp creates a new tar writer and starts upload to storage.
35
// Upload will block until the tar file is finished writing.
36
// If a name for the file is not given, default name is of
37
// the form `part_....tar.[Compressor file extension]`.
38
func (tarBall *StorageTarBall) SetUp(crypter crypto.Crypter, names ...string) {
39
	if tarBall.tarWriter == nil {
40
		if len(names) > 0 {
41
			tarBall.name = names[0]
42
		} else {
43
			tarBall.name = fmt.Sprintf("part_%0.3d.tar.%v", tarBall.partNumber, tarBall.uploader.Compression().FileExtension())
44
		}
45
		writeCloser := tarBall.startUpload(tarBall.name, crypter)
46

47
		tarBall.writeCloser = writeCloser
48
		tarBall.tarWriter = tar.NewWriter(writeCloser)
49
	}
50
}
51

52
// CloseTar closes the tar writer, flushing any unwritten data
53
// to the underlying writer before also closing the underlying writer.
54
func (tarBall *StorageTarBall) CloseTar() error {
55
	err := tarBall.tarWriter.Close()
56
	if err != nil {
57
		return errors.Wrap(err, "CloseTar: failed to close tar writer")
58
	}
59

60
	err = tarBall.writeCloser.Close()
61
	if err != nil {
62
		return errors.Wrap(err, "CloseTar: failed to close underlying writer")
63
	}
64
	tracelog.InfoLogger.Printf("Finished writing part %d.\n", tarBall.partNumber)
65
	return nil
66
}
67

68
func (tarBall *StorageTarBall) AwaitUploads() {
69
	tarBall.uploader.Finish()
70
	if tarBall.uploader.Failed() {
71
		tracelog.ErrorLogger.Fatal("Unable to complete uploads")
72
	}
73
}
74

75
// TODO : unit tests
76
// startUpload creates a compressing writer and runs upload in the background once
77
// a compressed tar member is finished writing.
78
func (tarBall *StorageTarBall) startUpload(name string, crypter crypto.Crypter) io.WriteCloser {
79
	pipeReader, pipeWriter := io.Pipe()
80
	uploader := tarBall.uploader
81

82
	path := tarBall.backupName + TarPartitionFolderName + name
83

84
	tracelog.InfoLogger.Printf("Starting part %d ...\n", tarBall.partNumber)
85

86
	go func() {
87
		err := uploader.Upload(context.Background(), path, pipeReader)
88
		if compressingError, ok := err.(CompressAndEncryptError); ok {
89
			tracelog.ErrorLogger.Printf("could not upload '%s' due to compression error\n%+v\n", path, compressingError)
90
		}
91
		if err != nil {
92
			tracelog.ErrorLogger.Printf("upload: could not upload '%s'\n", path)
93
			tracelog.ErrorLogger.Printf("%v\n", err)
94
			err = pipeReader.Close()
95
			tracelog.ErrorLogger.FatalfOnError("Failed to close pipe: %v", err)
96
			tracelog.ErrorLogger.Fatalf(
97
				"Unable to continue the backup process because of the loss of a part %d.\n",
98
				tarBall.partNumber)
99
		}
100
	}()
101

102
	var writerToCompress io.WriteCloser = pipeWriter
103

104
	if crypter != nil {
105
		encryptedWriter, err := crypter.Encrypt(pipeWriter)
106

107
		if err != nil {
108
			tracelog.ErrorLogger.Fatal("upload: encryption error ", err)
109
		}
110

111
		writerToCompress = &utility.CascadeWriteCloser{WriteCloser: encryptedWriter, Underlying: pipeWriter}
112
	}
113

114
	return &utility.CascadeWriteCloser{WriteCloser: uploader.Compression().NewWriter(writerToCompress),
115
		Underlying: writerToCompress}
116
}
117

118
// Size accumulated in this tarball
119
func (tarBall *StorageTarBall) Size() int64 { return atomic.LoadInt64(tarBall.partSize) }
120

121
// AddSize to total Size
122
func (tarBall *StorageTarBall) AddSize(i int64) { atomic.AddInt64(tarBall.partSize, i) }
123

124
func (tarBall *StorageTarBall) TarWriter() *tar.Writer { return tarBall.tarWriter }
125