Fix rare race in CopyFrom

This commit is contained in:
Jack Christensen 2022-05-24 08:26:37 -05:00
parent 55e0b4c30e
commit bfaea9e7ec

View File

@ -1184,8 +1184,11 @@ func (pgConn *PgConn) CopyFrom(ctx context.Context, r io.Reader, sql string) (Co
abortCopyChan := make(chan struct{}) abortCopyChan := make(chan struct{})
copyErrChan := make(chan error, 1) copyErrChan := make(chan error, 1)
signalMessageChan := pgConn.signalMessage() signalMessageChan := pgConn.signalMessage()
senderDoneChan := make(chan struct{})
go func() { go func() {
defer close(senderDoneChan)
buf := make([]byte, 0, 65536) buf := make([]byte, 0, 65536)
buf = append(buf, 'd') buf = append(buf, 'd')
sp := len(buf) sp := len(buf)
@ -1239,6 +1242,7 @@ func (pgConn *PgConn) CopyFrom(ctx context.Context, r io.Reader, sql string) (Co
} }
} }
close(abortCopyChan) close(abortCopyChan)
<-senderDoneChan
if copyErr == io.EOF || pgErr != nil { if copyErr == io.EOF || pgErr != nil {
pgConn.frontend.Send(&pgproto3.CopyDone{}) pgConn.frontend.Send(&pgproto3.CopyDone{})