storj/pkg/piecestore/rpc/server/readerwriter.go
Alexander Leitner b0db33f919
Bandwidth accounting (#134)
* captplanet standalone farmer setup

* Bandwidth Allocation

* utils.Close method changed to utils.LogClose

* Get build temporarily working

* Get/Put for PSClient should take payer bandwidth allocations rather than the NewPSClient function

* Update example client to reflect changes in client API

* Update ecclient to use latest PSClient, Make NewPSClient return error also

* Updated pieceranger tests to check for errors; sign method should take byte array

* Handle defers in store.go better

* Fix defer functions in psdb.go

* fun times

* Protobuf bandwidthallocation data is now a byte array

* Remove psservice package and merge it into pstore server

* Write wrapper for database calls

* Change all expiration names in protobuf to be more informative; add defer in retrieve; remove old comment

* Make PSDB tests implementation independent rather than method independent

* get rid of payer, renter in ecclient

* add context monitoring in store and retrieve
2018-08-17 13:40:15 -04:00

81 lines
2.0 KiB
Go

// Copyright (C) 2018 Storj Labs, Inc.
// See LICENSE for copying information.
package server
import (
"github.com/gogo/protobuf/proto"
"storj.io/storj/pkg/utils"
pb "storj.io/storj/protos/piecestore"
)
// StreamWriter -- Struct for writing piece to server upload stream
type StreamWriter struct {
server *Server
stream pb.PieceStoreRoutes_RetrieveServer
}
// NewStreamWriter returns a new StreamWriter
func NewStreamWriter(s *Server, stream pb.PieceStoreRoutes_RetrieveServer) *StreamWriter {
return &StreamWriter{server: s, stream: stream}
}
// Write -- Write method for piece upload to stream for Server.Retrieve
func (s *StreamWriter) Write(b []byte) (int, error) {
// Write the buffer to the stream we opened earlier
if err := s.stream.Send(&pb.PieceRetrievalStream{Size: int64(len(b)), Content: b}); err != nil {
return 0, err
}
return len(b), nil
}
// StreamReader is a struct for Retrieving data from server
type StreamReader struct {
src *utils.ReaderSource
bandwidthAllocation *pb.RenterBandwidthAllocation
currentTotal int64
}
// NewStreamReader returns a new StreamReader for Server.Store
func NewStreamReader(s *Server, stream pb.PieceStoreRoutes_StoreServer) *StreamReader {
sr := &StreamReader{}
sr.src = utils.NewReaderSource(func() ([]byte, error) {
recv, err := stream.Recv()
if err != nil {
return nil, err
}
pd := recv.GetPiecedata()
ba := recv.GetBandwidthallocation()
if ba != nil {
if err = s.verifySignature(ba); err != nil {
return nil, err
}
}
deserializedData := &pb.RenterBandwidthAllocation_Data{}
err = proto.Unmarshal(ba.GetData(), deserializedData)
if err != nil {
return nil, err
}
// Update bandwidthallocation to be stored
if deserializedData.GetTotal() > sr.currentTotal {
sr.bandwidthAllocation = ba
sr.currentTotal = deserializedData.GetTotal()
}
return pd.GetContent(), nil
})
return sr
}
// Read -- Read method for piece download from stream
func (s *StreamReader) Read(b []byte) (int, error) {
return s.src.Read(b)
}