forked from mongodb/mongo-go-driver
-
Notifications
You must be signed in to change notification settings - Fork 0
/
upload_stream.go
226 lines (189 loc) · 5.59 KB
/
upload_stream.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
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
// Copyright (C) MongoDB, Inc. 2017-present.
//
// Licensed under the Apache License, Version 2.0 (the "License"); you may
// not use this file except in compliance with the License. You may obtain
// a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
package gridfs
import (
"errors"
"context"
"time"
"math"
"go.mongodb.org/mongo-driver/bson/primitive"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/x/bsonx"
)
// UploadBufferSize is the size in bytes of one stream batch. Chunks will be written to the db after the sum of chunk
// lengths is equal to the batch size.
const UploadBufferSize = 16 * 1024 * 1024 // 16 MiB
// ErrStreamClosed is an error returned if an operation is attempted on a closed/aborted stream.
var ErrStreamClosed = errors.New("stream is closed or aborted")
// UploadStream is used to upload files in chunks.
type UploadStream struct {
*Upload // chunk size and metadata
FileID interface{}
chunkIndex int
chunksColl *mongo.Collection // collection to store file chunks
filename string
filesColl *mongo.Collection // collection to store file metadata
closed bool
buffer []byte
bufferIndex int
fileLen int64
writeDeadline time.Time
}
// NewUploadStream creates a new upload stream.
func newUploadStream(upload *Upload, fileID interface{}, filename string, chunks, files *mongo.Collection) *UploadStream {
return &UploadStream{
Upload: upload,
FileID: fileID,
chunksColl: chunks,
filename: filename,
filesColl: files,
buffer: make([]byte, UploadBufferSize),
}
}
// Close closes this upload stream.
func (us *UploadStream) Close() error {
if us.closed {
return ErrStreamClosed
}
ctx, cancel := deadlineContext(us.writeDeadline)
if cancel != nil {
defer cancel()
}
if us.bufferIndex != 0 {
if err := us.uploadChunks(ctx, true); err != nil {
return err
}
}
if err := us.createFilesCollDoc(ctx); err != nil {
return err
}
us.closed = true
return nil
}
// SetWriteDeadline sets the write deadline for this stream.
func (us *UploadStream) SetWriteDeadline(t time.Time) error {
if us.closed {
return ErrStreamClosed
}
us.writeDeadline = t
return nil
}
// Write transfers the contents of a byte slice into this upload stream. If the stream's underlying buffer fills up,
// the buffer will be uploaded as chunks to the server. Implements the io.Writer interface.
func (us *UploadStream) Write(p []byte) (int, error) {
if us.closed {
return 0, ErrStreamClosed
}
var ctx context.Context
ctx, cancel := deadlineContext(us.writeDeadline)
if cancel != nil {
defer cancel()
}
origLen := len(p)
for {
if len(p) == 0 {
break
}
n := copy(us.buffer[us.bufferIndex:], p) // copy as much as possible
p = p[n:]
us.bufferIndex += n
if us.bufferIndex == UploadBufferSize {
err := us.uploadChunks(ctx, false)
if err != nil {
return 0, err
}
}
}
return origLen, nil
}
// Abort closes the stream and deletes all file chunks that have already been written.
func (us *UploadStream) Abort() error {
if us.closed {
return ErrStreamClosed
}
ctx, cancel := deadlineContext(us.writeDeadline)
if cancel != nil {
defer cancel()
}
id, err := convertFileID(us.FileID)
if err != nil {
return err
}
_, err = us.chunksColl.DeleteMany(ctx, bsonx.Doc{{"files_id", id}})
if err != nil {
return err
}
us.closed = true
return nil
}
// uploadChunks uploads the current buffer as a series of chunks to the bucket
// if uploadPartial is true, any data at the end of the buffer that is smaller than a chunk will be uploaded as a partial
// chunk. if it is false, the data will be moved to the front of the buffer.
// uploadChunks sets us.bufferIndex to the next available index in the buffer after uploading
func (us *UploadStream) uploadChunks(ctx context.Context, uploadPartial bool) error {
chunks := float64(us.bufferIndex) / float64(us.chunkSize)
numChunks := int(math.Ceil(chunks))
if !uploadPartial {
numChunks = int(math.Floor(chunks))
}
docs := make([]interface{}, int(numChunks))
id, err := convertFileID(us.FileID)
if err != nil {
return err
}
begChunkIndex := us.chunkIndex
for i := 0; i < us.bufferIndex; i += int(us.chunkSize) {
endIndex := i + int(us.chunkSize)
if us.bufferIndex-i < int(us.chunkSize) {
// partial chunk
if !uploadPartial {
break
}
endIndex = us.bufferIndex
}
chunkData := us.buffer[i:endIndex]
docs[us.chunkIndex-begChunkIndex] = bsonx.Doc{
{"_id", bsonx.ObjectID(primitive.NewObjectID())},
{"files_id", id},
{"n", bsonx.Int32(int32(us.chunkIndex))},
{"data", bsonx.Binary(0x00, chunkData)},
}
us.chunkIndex++
us.fileLen += int64(len(chunkData))
}
_, err = us.chunksColl.InsertMany(ctx, docs)
if err != nil {
return err
}
// copy any remaining bytes to beginning of buffer and set buffer index
bytesUploaded := numChunks * int(us.chunkSize)
if bytesUploaded != UploadBufferSize && !uploadPartial {
copy(us.buffer[0:], us.buffer[bytesUploaded:us.bufferIndex])
}
us.bufferIndex = UploadBufferSize - bytesUploaded
return nil
}
func (us *UploadStream) createFilesCollDoc(ctx context.Context) error {
id, err := convertFileID(us.FileID)
if err != nil {
return err
}
doc := bsonx.Doc{
{"_id", id},
{"length", bsonx.Int64(us.fileLen)},
{"chunkSize", bsonx.Int32(us.chunkSize)},
{"uploadDate", bsonx.DateTime(time.Now().UnixNano() / int64(time.Millisecond))},
{"filename", bsonx.String(us.filename)},
}
if us.metadata != nil {
doc = append(doc, bsonx.Elem{"metadata", bsonx.Document(us.metadata)})
}
_, err = us.filesColl.InsertOne(ctx, doc)
if err != nil {
return err
}
return nil
}