-
Notifications
You must be signed in to change notification settings - Fork 3
/
writer.go
120 lines (98 loc) · 2.57 KB
/
writer.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
package seekable
import (
"fmt"
"io"
"sync"
"go.uber.org/multierr"
)
var (
_ io.Writer = (*WriterImpl)(nil)
_ io.Closer = (*WriterImpl)(nil)
)
// Environment can be used to inject a custom file reader that is different from normal ReadSeeker.
// This is useful when, for example there is a custom chunking code.
type WEnvironment interface {
// WriteFrame is called each time frame is encoded and needs to be written upstream.
WriteFrame(p []byte) (n int, err error)
// WriteSeekTable is called on Close to flush the seek table.
WriteSeekTable(p []byte) (n int, err error)
}
// writerEnvImpl is the environment implementation of for the underlying ReadSeeker.
type writerEnvImpl struct {
w io.Writer
}
func (w *writerEnvImpl) WriteFrame(p []byte) (n int, err error) {
return w.w.Write(p)
}
func (w *writerEnvImpl) WriteSeekTable(p []byte) (n int, err error) {
return w.w.Write(p)
}
type WriterImpl struct {
enc ZSTDEncoder
frameEntries []SeekTableEntry
o writerOptions
once *sync.Once
}
type Writer interface {
io.WriteCloser
}
type ZSTDEncoder interface {
EncodeAll(src, dst []byte) []byte
}
// NewWriter wraps the passed io.Writer and Encoder into and indexed ZSTD stream.
// Resulting stream then can be randomly accessed through the Reader and Decoder interfaces.
func NewWriter(w io.Writer, encoder ZSTDEncoder, opts ...WOption) (Writer, error) {
sw := WriterImpl{
once: &sync.Once{},
enc: encoder,
}
sw.o.setDefault()
for _, o := range opts {
err := o(&sw.o)
if err != nil {
return nil, err
}
}
if sw.o.env == nil {
sw.o.env = &writerEnvImpl{
w: w,
}
}
return &sw, nil
}
// Write writes a chunk of data as a separate frame into the datastream.
//
// Note that Write does not do any coalescing nor splitting of data,
// so each write will map to a separate ZSTD Frame.
func (s *WriterImpl) Write(src []byte) (int, error) {
dst, err := s.Encode(src)
if err != nil {
return 0, err
}
n, err := s.o.env.WriteFrame(dst)
if err != nil {
return 0, err
}
if n != len(dst) {
return 0, fmt.Errorf("partial write: %d out of %d", n, len(dst))
}
return len(src), nil
}
// Close implement io.Closer interface. It writes the seek table footer
// and releases occupied memory.
//
// Caller is still responsible to Close the underlying writer.
func (s *WriterImpl) Close() (err error) {
s.once.Do(func() {
err = multierr.Append(err, s.writeSeekTable())
})
return
}
func (s *WriterImpl) writeSeekTable() error {
seekTableBytes, err := s.EndStream()
if err != nil {
return err
}
_, err = s.o.env.WriteSeekTable(seekTableBytes)
return err
}