mirror of https://github.com/cbeuw/Cloak
Control flow optimisations
This commit is contained in:
parent
3ad04aa7e9
commit
cbd71fae6d
|
|
@ -14,7 +14,6 @@ import (
|
||||||
// it won't get chopped up into individual bytes
|
// it won't get chopped up into individual bytes
|
||||||
type datagramBufferedPipe struct {
|
type datagramBufferedPipe struct {
|
||||||
pLens []int
|
pLens []int
|
||||||
// lazily allocated
|
|
||||||
buf *bytes.Buffer
|
buf *bytes.Buffer
|
||||||
closed bool
|
closed bool
|
||||||
rwCond *sync.Cond
|
rwCond *sync.Cond
|
||||||
|
|
@ -27,6 +26,7 @@ type datagramBufferedPipe struct {
|
||||||
func NewDatagramBufferedPipe() *datagramBufferedPipe {
|
func NewDatagramBufferedPipe() *datagramBufferedPipe {
|
||||||
d := &datagramBufferedPipe{
|
d := &datagramBufferedPipe{
|
||||||
rwCond: sync.NewCond(&sync.Mutex{}),
|
rwCond: sync.NewCond(&sync.Mutex{}),
|
||||||
|
buf: new(bytes.Buffer),
|
||||||
}
|
}
|
||||||
return d
|
return d
|
||||||
}
|
}
|
||||||
|
|
@ -34,9 +34,6 @@ func NewDatagramBufferedPipe() *datagramBufferedPipe {
|
||||||
func (d *datagramBufferedPipe) Read(target []byte) (int, error) {
|
func (d *datagramBufferedPipe) Read(target []byte) (int, error) {
|
||||||
d.rwCond.L.Lock()
|
d.rwCond.L.Lock()
|
||||||
defer d.rwCond.L.Unlock()
|
defer d.rwCond.L.Unlock()
|
||||||
if d.buf == nil {
|
|
||||||
d.buf = new(bytes.Buffer)
|
|
||||||
}
|
|
||||||
for {
|
for {
|
||||||
if d.closed && len(d.pLens) == 0 {
|
if d.closed && len(d.pLens) == 0 {
|
||||||
return 0, io.EOF
|
return 0, io.EOF
|
||||||
|
|
@ -72,9 +69,6 @@ func (d *datagramBufferedPipe) Read(target []byte) (int, error) {
|
||||||
func (d *datagramBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
|
func (d *datagramBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
|
||||||
d.rwCond.L.Lock()
|
d.rwCond.L.Lock()
|
||||||
defer d.rwCond.L.Unlock()
|
defer d.rwCond.L.Unlock()
|
||||||
if d.buf == nil {
|
|
||||||
d.buf = new(bytes.Buffer)
|
|
||||||
}
|
|
||||||
for {
|
for {
|
||||||
if d.closed && len(d.pLens) == 0 {
|
if d.closed && len(d.pLens) == 0 {
|
||||||
return 0, io.EOF
|
return 0, io.EOF
|
||||||
|
|
@ -115,9 +109,6 @@ func (d *datagramBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
|
||||||
func (d *datagramBufferedPipe) Write(f *Frame) (toBeClosed bool, err error) {
|
func (d *datagramBufferedPipe) Write(f *Frame) (toBeClosed bool, err error) {
|
||||||
d.rwCond.L.Lock()
|
d.rwCond.L.Lock()
|
||||||
defer d.rwCond.L.Unlock()
|
defer d.rwCond.L.Unlock()
|
||||||
if d.buf == nil {
|
|
||||||
d.buf = new(bytes.Buffer)
|
|
||||||
}
|
|
||||||
for {
|
for {
|
||||||
if d.closed {
|
if d.closed {
|
||||||
return true, io.ErrClosedPipe
|
return true, io.ErrClosedPipe
|
||||||
|
|
|
||||||
|
|
@ -11,7 +11,6 @@ import (
|
||||||
|
|
||||||
// The point of a streamBufferedPipe is that Read() will block until data is available
|
// The point of a streamBufferedPipe is that Read() will block until data is available
|
||||||
type streamBufferedPipe struct {
|
type streamBufferedPipe struct {
|
||||||
// only alloc when on first Read or Write
|
|
||||||
buf *bytes.Buffer
|
buf *bytes.Buffer
|
||||||
|
|
||||||
closed bool
|
closed bool
|
||||||
|
|
@ -25,6 +24,7 @@ type streamBufferedPipe struct {
|
||||||
func NewStreamBufferedPipe() *streamBufferedPipe {
|
func NewStreamBufferedPipe() *streamBufferedPipe {
|
||||||
p := &streamBufferedPipe{
|
p := &streamBufferedPipe{
|
||||||
rwCond: sync.NewCond(&sync.Mutex{}),
|
rwCond: sync.NewCond(&sync.Mutex{}),
|
||||||
|
buf: new(bytes.Buffer),
|
||||||
}
|
}
|
||||||
return p
|
return p
|
||||||
}
|
}
|
||||||
|
|
@ -32,9 +32,6 @@ func NewStreamBufferedPipe() *streamBufferedPipe {
|
||||||
func (p *streamBufferedPipe) Read(target []byte) (int, error) {
|
func (p *streamBufferedPipe) Read(target []byte) (int, error) {
|
||||||
p.rwCond.L.Lock()
|
p.rwCond.L.Lock()
|
||||||
defer p.rwCond.L.Unlock()
|
defer p.rwCond.L.Unlock()
|
||||||
if p.buf == nil {
|
|
||||||
p.buf = new(bytes.Buffer)
|
|
||||||
}
|
|
||||||
for {
|
for {
|
||||||
if p.closed && p.buf.Len() == 0 {
|
if p.closed && p.buf.Len() == 0 {
|
||||||
return 0, io.EOF
|
return 0, io.EOF
|
||||||
|
|
@ -64,9 +61,6 @@ func (p *streamBufferedPipe) Read(target []byte) (int, error) {
|
||||||
func (p *streamBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
|
func (p *streamBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
|
||||||
p.rwCond.L.Lock()
|
p.rwCond.L.Lock()
|
||||||
defer p.rwCond.L.Unlock()
|
defer p.rwCond.L.Unlock()
|
||||||
if p.buf == nil {
|
|
||||||
p.buf = new(bytes.Buffer)
|
|
||||||
}
|
|
||||||
for {
|
for {
|
||||||
if p.closed && p.buf.Len() == 0 {
|
if p.closed && p.buf.Len() == 0 {
|
||||||
return 0, io.EOF
|
return 0, io.EOF
|
||||||
|
|
@ -104,9 +98,6 @@ func (p *streamBufferedPipe) WriteTo(w io.Writer) (n int64, err error) {
|
||||||
func (p *streamBufferedPipe) Write(input []byte) (int, error) {
|
func (p *streamBufferedPipe) Write(input []byte) (int, error) {
|
||||||
p.rwCond.L.Lock()
|
p.rwCond.L.Lock()
|
||||||
defer p.rwCond.L.Unlock()
|
defer p.rwCond.L.Unlock()
|
||||||
if p.buf == nil {
|
|
||||||
p.buf = new(bytes.Buffer)
|
|
||||||
}
|
|
||||||
for {
|
for {
|
||||||
if p.closed {
|
if p.closed {
|
||||||
return 0, io.ErrClosedPipe
|
return 0, io.ErrClosedPipe
|
||||||
|
|
|
||||||
|
|
@ -72,8 +72,35 @@ func (sb *switchboard) addConn(conn net.Conn) {
|
||||||
|
|
||||||
// a pointer to connId is passed here so that the switchboard can reassign it if that connId isn't usable
|
// a pointer to connId is passed here so that the switchboard can reassign it if that connId isn't usable
|
||||||
func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
|
func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
|
||||||
writeAndRegUsage := func(conn net.Conn, d []byte) (int, error) {
|
sb.valve.txWait(len(data))
|
||||||
n, err = conn.Write(d)
|
if atomic.LoadUint32(&sb.broken) == 1 || sb.connsCount() == 0 {
|
||||||
|
return 0, errBrokenSwitchboard
|
||||||
|
}
|
||||||
|
|
||||||
|
var conn net.Conn
|
||||||
|
switch sb.strategy {
|
||||||
|
case UNIFORM_SPREAD:
|
||||||
|
_, conn, err = sb.pickRandConn()
|
||||||
|
if err != nil {
|
||||||
|
return 0, errBrokenSwitchboard
|
||||||
|
}
|
||||||
|
case FIXED_CONN_MAPPING:
|
||||||
|
connI, ok := sb.conns.Load(*connId)
|
||||||
|
if ok {
|
||||||
|
conn = connI.(net.Conn)
|
||||||
|
} else {
|
||||||
|
var newConnId uint32
|
||||||
|
newConnId, conn, err = sb.pickRandConn()
|
||||||
|
if err != nil {
|
||||||
|
return 0, errBrokenSwitchboard
|
||||||
|
}
|
||||||
|
*connId = newConnId
|
||||||
|
}
|
||||||
|
default:
|
||||||
|
return 0, errors.New("unsupported traffic distribution strategy")
|
||||||
|
}
|
||||||
|
|
||||||
|
n, err = conn.Write(data)
|
||||||
if err != nil {
|
if err != nil {
|
||||||
sb.conns.Delete(*connId)
|
sb.conns.Delete(*connId)
|
||||||
sb.session.SetTerminalMsg("failed to write to remote " + err.Error())
|
sb.session.SetTerminalMsg("failed to write to remote " + err.Error())
|
||||||
|
|
@ -84,36 +111,6 @@ func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
|
||||||
return n, nil
|
return n, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
sb.valve.txWait(len(data))
|
|
||||||
if atomic.LoadUint32(&sb.broken) == 1 || sb.connsCount() == 0 {
|
|
||||||
return 0, errBrokenSwitchboard
|
|
||||||
}
|
|
||||||
|
|
||||||
switch sb.strategy {
|
|
||||||
case UNIFORM_SPREAD:
|
|
||||||
_, conn, err := sb.pickRandConn()
|
|
||||||
if err != nil {
|
|
||||||
return 0, errBrokenSwitchboard
|
|
||||||
}
|
|
||||||
return writeAndRegUsage(conn, data)
|
|
||||||
case FIXED_CONN_MAPPING:
|
|
||||||
connI, ok := sb.conns.Load(*connId)
|
|
||||||
if ok {
|
|
||||||
conn := connI.(net.Conn)
|
|
||||||
return writeAndRegUsage(conn, data)
|
|
||||||
} else {
|
|
||||||
newConnId, conn, err := sb.pickRandConn()
|
|
||||||
if err != nil {
|
|
||||||
return 0, errBrokenSwitchboard
|
|
||||||
}
|
|
||||||
*connId = newConnId
|
|
||||||
return writeAndRegUsage(conn, data)
|
|
||||||
}
|
|
||||||
default:
|
|
||||||
return 0, errors.New("unsupported traffic distribution strategy")
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
// returns a random connId
|
// returns a random connId
|
||||||
func (sb *switchboard) pickRandConn() (uint32, net.Conn, error) {
|
func (sb *switchboard) pickRandConn() (uint32, net.Conn, error) {
|
||||||
connCount := sb.connsCount()
|
connCount := sb.connsCount()
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue