fluent: check if context is cancelled before performing I/O on the connection

There are a number of places in the code where we do not check if the context
is cancelled before reading from or writing to the connection. This commit adds
those checks.

Signed-off-by: Anirudh Aithal <[email protected]>

cr: https://code.amazon.com/reviews/CR-194284782

Author: aaithal

fluent/fluent.go

@@ -594,10 +594,15 @@ func (f *Fluent) writeWithRetry(ctx context.Context, msg *msgToSend) error {
 	return fmt.Errorf("fluent#write: failed to write after %d attempts", f.Config.MaxRetry)
 }
 
-func (f *Fluent) syncWriteMessage(msg *msgToSend) error {
+func (f *Fluent) syncWriteMessage(ctx context.Context, msg *msgToSend) error {
 	f.muconn.RLock()
 	defer f.muconn.RUnlock()
 
+	// Check if context is cancelled. If it is, we can return early here.
+	if err := ctx.Err(); err != nil {
+		return errIsClosing
+	}
+
 	if f.conn == nil {
 		return fmt.Errorf("connection has been closed before writing to it")
 	}
@@ -613,12 +618,18 @@ func (f *Fluent) syncWriteMessage(msg *msgToSend) error {
 	return err
 }
 
-func (f *Fluent) syncReadAck() (*AckResp, error) {
+func (f *Fluent) syncReadAck(ctx context.Context) (*AckResp, error) {
 	f.muconn.RLock()
 	defer f.muconn.RUnlock()
 
 	resp := &AckResp{}
 	var err error
+
+	// Check if context is cancelled. If it is, we can return early here.
+	if err := ctx.Err(); err != nil {
+		return resp, errIsClosing
+	}
+
 	if f.Config.MarshalAsJSON {
 		dec := json.NewDecoder(f.conn)
 		err = dec.Decode(resp)
@@ -642,14 +653,14 @@ func (f *Fluent) write(ctx context.Context, msg *msgToSend) (bool, error) {
 		return false, fmt.Errorf("fluent#write: %v", err)
 	}
 
-	if err := f.syncWriteMessage(msg); err != nil {
+	if err := f.syncWriteMessage(ctx, msg); err != nil {
 		f.syncClose(false)
 		return true, fmt.Errorf("fluent#write: %v", err)
 	}
 
 	// Acknowledgment check
 	if msg.ack != "" {
-		resp, err := f.syncReadAck()
+		resp, err := f.syncReadAck(ctx)
 		if err != nil {
 			fmt.Fprintf(os.Stderr, "fluent#write: error reading message response ack %v", err)
 			f.syncClose(false)