make putting back packet buffers a member function of packetBuffer

This commit is contained in:
Marten Seemann 2018-12-30 15:28:52 +07:00
parent d3ea7c0c4c
commit b32fb438af
5 changed files with 35 additions and 32 deletions

View file

@ -22,6 +22,23 @@ func (b *packetBuffer) Split() {
b.refCount++ b.refCount++
} }
// Release decreases the refCount.
// It should be called when processing the packet is finished.
// When the refCount reaches 0, the packet buffer is put back into the pool.
func (b *packetBuffer) Release() {
if cap(b.Slice) != int(protocol.MaxReceivePacketSize) {
panic("putPacketBuffer called with packet of wrong size!")
}
b.refCount--
if b.refCount < 0 {
panic("negative packetBuffer refCount")
}
// only put the packetBuffer back if it's not used any more
if b.refCount == 0 {
bufferPool.Put(b)
}
}
var bufferPool sync.Pool var bufferPool sync.Pool
func getPacketBuffer() *packetBuffer { func getPacketBuffer() *packetBuffer {
@ -31,20 +48,6 @@ func getPacketBuffer() *packetBuffer {
return buf return buf
} }
func putPacketBuffer(buf *packetBuffer) {
if cap(buf.Slice) != int(protocol.MaxReceivePacketSize) {
panic("putPacketBuffer called with packet of wrong size!")
}
buf.refCount--
if buf.refCount < 0 {
panic("negative packetBuffer refCount")
}
// only put the packetBuffer back if it's not used any more
if buf.refCount == 0 {
bufferPool.Put(buf)
}
}
func init() { func init() {
bufferPool.New = func() interface{} { bufferPool.New = func() interface{} {
return &packetBuffer{ return &packetBuffer{

View file

@ -13,31 +13,31 @@ var _ = Describe("Buffer Pool", func() {
Expect(buf.Slice).To(HaveCap(int(protocol.MaxReceivePacketSize))) Expect(buf.Slice).To(HaveCap(int(protocol.MaxReceivePacketSize)))
}) })
It("puts buffers back", func() { It("releases buffers", func() {
buf := getPacketBuffer() buf := getPacketBuffer()
putPacketBuffer(buf) buf.Release()
}) })
It("panics if wrong-sized buffers are passed", func() { It("panics if wrong-sized buffers are passed", func() {
buf := getPacketBuffer() buf := getPacketBuffer()
buf.Slice = make([]byte, 10) buf.Slice = make([]byte, 10)
Expect(func() { putPacketBuffer(buf) }).To(Panic()) Expect(func() { buf.Release() }).To(Panic())
}) })
It("panics if it is put pack twice", func() { It("panics if it is released twice", func() {
buf := getPacketBuffer() buf := getPacketBuffer()
putPacketBuffer(buf) buf.Release()
Expect(func() { putPacketBuffer(buf) }).To(Panic()) Expect(func() { buf.Release() }).To(Panic())
}) })
It("waits until all parts have been put back", func() { It("waits until all parts have been released", func() {
buf := getPacketBuffer() buf := getPacketBuffer()
buf.Split() buf.Split()
buf.Split() buf.Split()
// now we have 3 parts // now we have 3 parts
putPacketBuffer(buf) buf.Release()
putPacketBuffer(buf) buf.Release()
putPacketBuffer(buf) buf.Release()
Expect(func() { putPacketBuffer(buf) }).To(Panic()) Expect(func() { buf.Release() }).To(Panic())
}) })
}) })

View file

@ -169,7 +169,7 @@ func (h *packetHandlerMap) handlePacket(
// We still need to process the packets that were successfully parsed before. // We still need to process the packets that were successfully parsed before.
} }
if len(packets) == 0 { if len(packets) == 0 {
putPacketBuffer(buffer) buffer.Release()
return return
} }
h.handleParsedPackets(packets) h.handleParsedPackets(packets)

View file

@ -321,20 +321,20 @@ func (s *server) handlePacket(p *receivedPacket) {
return return
} }
putPacketBuffer(p.buffer)
// TODO(#943): send Stateless Reset // TODO(#943): send Stateless Reset
p.buffer.Release()
} }
func (s *server) handleInitial(p *receivedPacket) { func (s *server) handleInitial(p *receivedPacket) {
s.logger.Debugf("<- Received Initial packet.") s.logger.Debugf("<- Received Initial packet.")
sess, connID, err := s.handleInitialImpl(p) sess, connID, err := s.handleInitialImpl(p)
if err != nil { if err != nil {
putPacketBuffer(p.buffer) p.buffer.Release()
s.logger.Errorf("Error occurred handling initial packet: %s", err) s.logger.Errorf("Error occurred handling initial packet: %s", err)
return return
} }
if sess == nil { // a retry was done if sess == nil { // a retry was done
putPacketBuffer(p.buffer) p.buffer.Release()
return return
} }
// Don't put the packet buffer back if a new session was created. // Don't put the packet buffer back if a new session was created.
@ -461,7 +461,7 @@ func (s *server) sendRetry(remoteAddr net.Addr, hdr *wire.Header) error {
} }
func (s *server) sendVersionNegotiationPacket(p *receivedPacket) { func (s *server) sendVersionNegotiationPacket(p *receivedPacket) {
defer putPacketBuffer(p.buffer) defer p.buffer.Release()
hdr := p.hdr hdr := p.hdr
s.logger.Debugf("Client offered version %s, sending Version Negotiation", hdr.Version) s.logger.Debugf("Client offered version %s, sending Version Negotiation", hdr.Version)
data, err := wire.ComposeVersionNegotiation(hdr.SrcConnectionID, hdr.DestConnectionID, s.config.Versions) data, err := wire.ComposeVersionNegotiation(hdr.SrcConnectionID, hdr.DestConnectionID, s.config.Versions)

View file

@ -479,7 +479,7 @@ func (s *session) handlePacketImpl(p *receivedPacket) bool /* was the packet suc
defer func() { defer func() {
// Put back the packet buffer if the packet wasn't queued for later decryption. // Put back the packet buffer if the packet wasn't queued for later decryption.
if !wasQueued { if !wasQueued {
putPacketBuffer(p.buffer) p.buffer.Release()
} }
}() }()
@ -962,7 +962,7 @@ func (s *session) sendPacket() (bool, error) {
} }
func (s *session) sendPackedPacket(packet *packedPacket) error { func (s *session) sendPackedPacket(packet *packedPacket) error {
defer putPacketBuffer(packet.buffer) defer packet.buffer.Release()
s.logPacket(packet) s.logPacket(packet)
return s.conn.Write(packet.raw) return s.conn.Write(packet.raw)
} }