Introduce a postWrite hook.

If you are sending large amounts of data (multiple megabytes) and you don't want to allocate that large a byte slice, one option is to send the data and have a mechanism to Wake() the connection up after the write is done. This will require a postWrite(conn) hook that also provides a reference to the connection for which data was just written. This way the application can maintain state in the context and in the postWrite hook, call Wake() and send the next batch of data when the Data() hook gets called.
This commit is contained in:
Karthik Abram 2023-07-07 15:31:41 -04:00
parent fe60817601
commit e1e8a34ba8
3 changed files with 17 additions and 0 deletions

View File

@ -115,6 +115,8 @@ type Events struct {
Detached func(c Conn, rwc io.ReadWriteCloser) (action Action)
// PreWrite fires just before any data is written to any client socket.
PreWrite func()
// PostWrite fires just after any data is written to any client socket.
PostWrite func(c Conn)
// Data fires when a connection sends the server data.
// The in parameter is the incoming data.
// Use the out return value to write data to the connection.

View File

@ -384,6 +384,9 @@ func stdloopRead(s *stdserver, l *stdloop, c *stdconn, in []byte) error {
s.events.PreWrite()
}
c.conn.Write(out)
if s.events.PostWrite != nil {
s.events.PostWrite(c)
}
}
switch action {
case Shutdown:
@ -405,6 +408,9 @@ func stdloopReadUDP(s *stdserver, l *stdloop, c *stdudpconn) error {
s.events.PreWrite()
}
s.lns[c.addrIndex].pconn.WriteTo(out, c.remoteAddr)
if s.events.PostWrite != nil {
s.events.PostWrite(c)
}
}
switch action {
case Shutdown:
@ -439,6 +445,9 @@ func stdloopAccept(s *stdserver, l *stdloop, c *stdconn) error {
s.events.PreWrite()
}
c.conn.Write(out)
if s.events.PostWrite != nil {
s.events.PostWrite(c)
}
}
if opts.TCPKeepAlive > 0 {
if c, ok := c.conn.(*net.TCPConn); ok {

View File

@ -333,6 +333,9 @@ func loopUDPRead(s *server, l *loop, lnidx, fd int) error {
s.events.PreWrite()
}
syscall.Sendto(fd, out, 0, sa)
if s.events.PostWrite != nil {
s.events.PostWrite(c)
}
}
switch action {
case Shutdown:
@ -391,6 +394,9 @@ func loopWrite(s *server, l *loop, c *conn) error {
if len(c.out) == 0 && c.action == None {
l.poll.ModRead(c.fd)
}
if len(c.out) == 0 && s.events.PostWrite != nil {
s.events.PostWrite(c)
}
return nil
}