a20a7db793
It provides an abstraction around the rpc details so that one can use dprc or gprc with the same code. It subsumes using the protobuf package directly for client interfaces as well as the pkg/transport package to perform dials. Change-Id: I8f5688bd71be8b0c766f13029128a77e5d46320b
79 lines
1.9 KiB
Go
79 lines
1.9 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/internal/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
|
|
timeout time.Duration
|
|
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 setting the timeout and sleeping to ensure the rate.
|
|
func (t *timedConn) Read(p []byte) (int, error) {
|
|
if t.timeout > 0 {
|
|
if err := t.SetReadDeadline(time.Now().Add(t.timeout)); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
|
|
start := t.now()
|
|
n, err := t.Conn.Read(p)
|
|
t.delay(start, n)
|
|
return n, err
|
|
}
|
|
|
|
// Write wraps the connection write setting the timeout and sleeping to ensure the rate.
|
|
func (t *timedConn) Write(p []byte) (int, error) {
|
|
if t.timeout > 0 {
|
|
if err := t.SetWriteDeadline(time.Now().Add(t.timeout)); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
|
|
start := t.now()
|
|
n, err := t.Conn.Write(p)
|
|
t.delay(start, n)
|
|
return n, err
|
|
}
|