storj/pkg/rpc/common.go
2019-11-14 21:46:15 +02:00

66 lines
1.6 KiB
Go

// Copyright (C) 2019 Storj Labs, Inc.
// See LICENSE for copying information.
package rpc
import (
"net"
"time"
"github.com/zeebo/errs"
"gopkg.in/spacemonkeygo/monkit.v2"
"storj.io/storj/private/memory"
)
//go:generate go run gen.go ../pb drpc compat_drpc.go
//go:generate go run gen.go ../pb grpc compat_grpc.go
var mon = monkit.Package()
// Error wraps all of the errors returned by this package.
var Error = errs.Class("rpccompat")
// timedConn wraps a net.Conn so that all reads and writes get the specified timeout and
// return bytes no faster than the rate. If the timeout or rate are zero, they are
// ignored.
type timedConn struct {
net.Conn
rate memory.Size
}
// now returns time.Now if there's a nonzero rate.
func (t *timedConn) now() (now time.Time) {
if t.rate > 0 {
now = time.Now()
}
return now
}
// delay ensures that we sleep to keep the rate if it is nonzero. n is the number of
// bytes in the read or write operation we need to delay.
func (t *timedConn) delay(start time.Time, n int) {
if t.rate > 0 {
expected := time.Duration(n * int(time.Second) / t.rate.Int())
if actual := time.Since(start); expected > actual {
time.Sleep(expected - actual)
}
}
}
// Read wraps the connection read and adds sleeping to ensure the rate.
func (t *timedConn) Read(p []byte) (int, error) {
start := t.now()
n, err := t.Conn.Read(p)
t.delay(start, n)
return n, err
}
// Write wraps the connection write and adds sleeping to ensure the rate.
func (t *timedConn) Write(p []byte) (int, error) {
start := t.now()
n, err := t.Conn.Write(p)
t.delay(start, n)
return n, err
}