5
0
mirror of https://github.com/cwinfo/yggdrasil-go.git synced 2024-11-26 13:01:39 +00:00

add Conn.ReadNoCopy and Conn.WriteNoCopy that transfer ownership of a slice instead of copying, have Read and Write use the NoCopy versions under the hood and just manage copying as needed

This commit is contained in:
Arceliar 2019-08-04 15:53:34 -05:00
parent 07f14f92ed
commit 5d5486049b

View File

@ -149,73 +149,90 @@ func (c *Conn) getDeadlineCancellation(value *atomic.Value) util.Cancellation {
} }
} }
func (c *Conn) Read(b []byte) (int, error) { // Used internally by Read, the caller is responsible for util.PutBytes when they're done.
// Take a copy of the session object func (c *Conn) ReadNoCopy() ([]byte, error) {
sinfo := c.session
cancel := c.getDeadlineCancellation(&c.readDeadline) cancel := c.getDeadlineCancellation(&c.readDeadline)
defer cancel.Cancel(nil) defer cancel.Cancel(nil)
// Wait for some traffic to come through from the session // Wait for some traffic to come through from the session
select { select {
case <-cancel.Finished(): case <-cancel.Finished():
if cancel.Error() == util.CancellationTimeoutError { if cancel.Error() == util.CancellationTimeoutError {
return 0, ConnError{errors.New("read timeout"), true, false, false, 0} return nil, ConnError{errors.New("read timeout"), true, false, false, 0}
} else { } else {
return 0, ConnError{errors.New("session closed"), false, false, true, 0} return nil, ConnError{errors.New("session closed"), false, false, true, 0}
} }
case bs := <-sinfo.recv: case bs := <-c.session.recv:
var err error return bs, nil
n := len(bs)
if len(bs) > len(b) {
n = len(b)
err = ConnError{errors.New("read buffer too small for entire packet"), false, true, false, 0}
}
// Copy results to the output slice and clean up
copy(b, bs)
util.PutBytes(bs)
// If we've reached this point then everything went to plan, return the
// number of bytes we populated back into the given slice
return n, err
} }
} }
func (c *Conn) Write(b []byte) (bytesWritten int, err error) { // Implements net.Conn.Read
sinfo := c.session func (c *Conn) Read(b []byte) (int, error) {
written := len(b) bs, err := c.ReadNoCopy()
if err != nil {
return 0, err
}
n := len(bs)
if len(bs) > len(b) {
n = len(b)
err = ConnError{errors.New("read buffer too small for entire packet"), false, true, false, 0}
}
// Copy results to the output slice and clean up
copy(b, bs)
util.PutBytes(bs)
// Return the number of bytes copied to the slice, along with any error
return n, err
}
// Used internally by Write, the caller must not reuse the argument bytes when no error occurs
func (c *Conn) WriteNoCopy(bs []byte) error {
var err error
sessionFunc := func() { sessionFunc := func() {
// Does the packet exceed the permitted size for the session? // Does the packet exceed the permitted size for the session?
if uint16(len(b)) > sinfo.getMTU() { if uint16(len(bs)) > c.session.getMTU() {
written, err = 0, ConnError{errors.New("packet too big"), true, false, false, int(sinfo.getMTU())} err = ConnError{errors.New("packet too big"), true, false, false, int(c.session.getMTU())}
return return
} }
// The rest of this work is session keep-alive traffic // The rest of this work is session keep-alive traffic
switch { switch {
case time.Since(sinfo.time) > 6*time.Second: case time.Since(c.session.time) > 6*time.Second:
if sinfo.time.Before(sinfo.pingTime) && time.Since(sinfo.pingTime) > 6*time.Second { if c.session.time.Before(c.session.pingTime) && time.Since(c.session.pingTime) > 6*time.Second {
// TODO double check that the above condition is correct // TODO double check that the above condition is correct
c.doSearch() c.doSearch()
} else { } else {
sinfo.core.sessions.ping(sinfo) c.core.sessions.ping(c.session)
} }
case sinfo.reset && sinfo.pingTime.Before(sinfo.time): case c.session.reset && c.session.pingTime.Before(c.session.time):
sinfo.core.sessions.ping(sinfo) c.core.sessions.ping(c.session)
default: // Don't do anything, to keep traffic throttled default: // Don't do anything, to keep traffic throttled
} }
} }
sinfo.doFunc(sessionFunc) c.session.doFunc(sessionFunc)
if written > 0 { if err == nil {
bs := append(util.GetBytes(), b...)
cancel := c.getDeadlineCancellation(&c.writeDeadline) cancel := c.getDeadlineCancellation(&c.writeDeadline)
defer cancel.Cancel(nil) defer cancel.Cancel(nil)
select { select {
case <-cancel.Finished(): case <-cancel.Finished():
if cancel.Error() == util.CancellationTimeoutError { if cancel.Error() == util.CancellationTimeoutError {
return 0, ConnError{errors.New("write timeout"), true, false, false, 0} err = ConnError{errors.New("write timeout"), true, false, false, 0}
} else { } else {
return 0, ConnError{errors.New("session closed"), false, false, true, 0} err = ConnError{errors.New("session closed"), false, false, true, 0}
} }
case sinfo.send <- bs: case c.session.send <- bs:
} }
} }
return err
}
// Implements net.Conn.Write
func (c *Conn) Write(b []byte) (int, error) {
written := len(b)
bs := append(util.GetBytes(), b...)
err := c.WriteNoCopy(bs)
if err != nil {
util.PutBytes(bs)
written = 0
}
return written, err return written, err
} }