chrislu
3 weeks ago
2 changed files with 312 additions and 132 deletions
@ -0,0 +1,214 @@ |
|||
package operation |
|||
|
|||
import ( |
|||
"bytes" |
|||
"context" |
|||
"crypto/md5" |
|||
"fmt" |
|||
"hash" |
|||
"io" |
|||
"sync" |
|||
|
|||
"github.com/seaweedfs/seaweedfs/weed/glog" |
|||
"github.com/seaweedfs/seaweedfs/weed/pb/filer_pb" |
|||
"github.com/seaweedfs/seaweedfs/weed/security" |
|||
) |
|||
|
|||
// ChunkedUploadResult contains the result of a chunked upload
|
|||
type ChunkedUploadResult struct { |
|||
FileChunks []*filer_pb.FileChunk |
|||
Md5Hash hash.Hash |
|||
TotalSize int64 |
|||
SmallContent []byte // For files smaller than threshold
|
|||
} |
|||
|
|||
// ChunkedUploadOption contains options for chunked uploads
|
|||
type ChunkedUploadOption struct { |
|||
ChunkSize int32 |
|||
SmallFileLimit int64 |
|||
Collection string |
|||
Replication string |
|||
DataCenter string |
|||
SaveSmallInline bool |
|||
Jwt security.EncodedJwt |
|||
MimeType string |
|||
AssignFunc func(ctx context.Context, count int) (*VolumeAssignRequest, *AssignResult, error) |
|||
} |
|||
|
|||
var chunkBufferPool = sync.Pool{ |
|||
New: func() interface{} { |
|||
return new(bytes.Buffer) |
|||
}, |
|||
} |
|||
|
|||
// UploadReaderInChunks reads from reader and uploads in chunks to volume servers
|
|||
// This prevents OOM by processing the stream in fixed-size chunks
|
|||
// Returns file chunks, MD5 hash, total size, and any small content stored inline
|
|||
func UploadReaderInChunks(ctx context.Context, reader io.Reader, opt *ChunkedUploadOption) (*ChunkedUploadResult, error) { |
|||
|
|||
md5Hash := md5.New() |
|||
var partReader = io.TeeReader(reader, md5Hash) |
|||
|
|||
var fileChunks []*filer_pb.FileChunk |
|||
var fileChunksLock sync.Mutex |
|||
var uploadErr error |
|||
var uploadErrLock sync.Mutex |
|||
var chunkOffset int64 = 0 |
|||
|
|||
var wg sync.WaitGroup |
|||
var bytesBufferCounter int64 = 4 |
|||
bytesBufferLimitChan := make(chan struct{}, bytesBufferCounter) |
|||
|
|||
for { |
|||
// Throttle buffer usage
|
|||
bytesBufferLimitChan <- struct{}{} |
|||
|
|||
// Check for errors from parallel uploads
|
|||
uploadErrLock.Lock() |
|||
if uploadErr != nil { |
|||
<-bytesBufferLimitChan |
|||
uploadErrLock.Unlock() |
|||
break |
|||
} |
|||
uploadErrLock.Unlock() |
|||
|
|||
// Get buffer from pool
|
|||
bytesBuffer := chunkBufferPool.Get().(*bytes.Buffer) |
|||
limitedReader := io.LimitReader(partReader, int64(opt.ChunkSize)) |
|||
bytesBuffer.Reset() |
|||
|
|||
// Read one chunk
|
|||
dataSize, err := bytesBuffer.ReadFrom(limitedReader) |
|||
if err != nil || dataSize == 0 { |
|||
chunkBufferPool.Put(bytesBuffer) |
|||
<-bytesBufferLimitChan |
|||
if err != nil { |
|||
uploadErrLock.Lock() |
|||
if uploadErr == nil { |
|||
uploadErr = err |
|||
} |
|||
uploadErrLock.Unlock() |
|||
} |
|||
break |
|||
} |
|||
|
|||
// For small files at offset 0, store inline instead of uploading
|
|||
if chunkOffset == 0 && opt.SaveSmallInline && dataSize < opt.SmallFileLimit { |
|||
smallContent := make([]byte, dataSize) |
|||
bytesBuffer.Read(smallContent) |
|||
chunkBufferPool.Put(bytesBuffer) |
|||
<-bytesBufferLimitChan |
|||
|
|||
return &ChunkedUploadResult{ |
|||
FileChunks: nil, |
|||
Md5Hash: md5Hash, |
|||
TotalSize: dataSize, |
|||
SmallContent: smallContent, |
|||
}, nil |
|||
} |
|||
|
|||
// Upload chunk in parallel goroutine
|
|||
wg.Add(1) |
|||
go func(offset int64, buf *bytes.Buffer) { |
|||
defer func() { |
|||
chunkBufferPool.Put(buf) |
|||
<-bytesBufferLimitChan |
|||
wg.Done() |
|||
}() |
|||
|
|||
// Assign volume for this chunk
|
|||
_, assignResult, assignErr := opt.AssignFunc(ctx, 1) |
|||
if assignErr != nil { |
|||
uploadErrLock.Lock() |
|||
if uploadErr == nil { |
|||
uploadErr = fmt.Errorf("assign volume: %w", assignErr) |
|||
} |
|||
uploadErrLock.Unlock() |
|||
return |
|||
} |
|||
|
|||
// Upload chunk data
|
|||
uploadUrl := fmt.Sprintf("http://%s/%s", assignResult.Url, assignResult.Fid) |
|||
uploadOption := &UploadOption{ |
|||
UploadUrl: uploadUrl, |
|||
Cipher: false, |
|||
IsInputCompressed: false, |
|||
MimeType: opt.MimeType, |
|||
PairMap: nil, |
|||
Jwt: opt.Jwt, |
|||
} |
|||
|
|||
uploader, uploaderErr := NewUploader() |
|||
if uploaderErr != nil { |
|||
uploadErrLock.Lock() |
|||
if uploadErr == nil { |
|||
uploadErr = fmt.Errorf("create uploader: %w", uploaderErr) |
|||
} |
|||
uploadErrLock.Unlock() |
|||
return |
|||
} |
|||
|
|||
uploadResult, uploadResultErr := uploader.UploadData(ctx, buf.Bytes(), uploadOption) |
|||
if uploadResultErr != nil { |
|||
uploadErrLock.Lock() |
|||
if uploadErr == nil { |
|||
uploadErr = fmt.Errorf("upload chunk: %w", uploadResultErr) |
|||
} |
|||
uploadErrLock.Unlock() |
|||
return |
|||
} |
|||
|
|||
// Create chunk entry
|
|||
fid, _ := filer_pb.ToFileIdObject(assignResult.Fid) |
|||
chunk := &filer_pb.FileChunk{ |
|||
FileId: assignResult.Fid, |
|||
Offset: offset, |
|||
Size: uint64(uploadResult.Size), |
|||
ETag: uploadResult.ContentMd5, |
|||
Fid: fid, |
|||
CipherKey: uploadResult.CipherKey, |
|||
} |
|||
|
|||
fileChunksLock.Lock() |
|||
fileChunks = append(fileChunks, chunk) |
|||
glog.V(4).Infof("uploaded chunk %d to %s [%d,%d)", len(fileChunks), chunk.FileId, offset, offset+int64(chunk.Size)) |
|||
fileChunksLock.Unlock() |
|||
|
|||
}(chunkOffset, bytesBuffer) |
|||
|
|||
// Update offset for next chunk
|
|||
chunkOffset += dataSize |
|||
|
|||
// If this was a partial chunk, we're done
|
|||
if dataSize < int64(opt.ChunkSize) { |
|||
break |
|||
} |
|||
} |
|||
|
|||
// Wait for all uploads to complete
|
|||
wg.Wait() |
|||
|
|||
// Check for errors
|
|||
if uploadErr != nil { |
|||
glog.Errorf("chunked upload failed: %v", uploadErr) |
|||
return nil, uploadErr |
|||
} |
|||
|
|||
// Sort chunks by offset
|
|||
// Note: We could use slices.SortFunc here, but keeping it simple for Go 1.20 compatibility
|
|||
for i := 0; i < len(fileChunks); i++ { |
|||
for j := i + 1; j < len(fileChunks); j++ { |
|||
if fileChunks[i].Offset > fileChunks[j].Offset { |
|||
fileChunks[i], fileChunks[j] = fileChunks[j], fileChunks[i] |
|||
} |
|||
} |
|||
} |
|||
|
|||
return &ChunkedUploadResult{ |
|||
FileChunks: fileChunks, |
|||
Md5Hash: md5Hash, |
|||
TotalSize: chunkOffset, |
|||
SmallContent: nil, |
|||
}, nil |
|||
} |
|||
|
|||
Write
Preview
Loading…
Cancel
Save
Reference in new issue