github.com/muhammadn/cortex@v1.9.1-0.20220510110439-46bb7000d03d/pkg/ingester/client/cortex_util.go (about)

     1  package client
     2  
     3  import (
     4  	context "context"
     5  )
     6  
     7  // SendQueryStream wraps the stream's Send() checking if the context is done
     8  // before calling Send().
     9  func SendQueryStream(s Ingester_QueryStreamServer, m *QueryStreamResponse) error {
    10  	return sendWithContextErrChecking(s.Context(), func() error {
    11  		return s.Send(m)
    12  	})
    13  }
    14  
    15  // SendTimeSeriesChunk wraps the stream's Send() checking if the context is done
    16  // before calling Send().
    17  func SendTimeSeriesChunk(s Ingester_TransferChunksClient, m *TimeSeriesChunk) error {
    18  	return sendWithContextErrChecking(s.Context(), func() error {
    19  		return s.Send(m)
    20  	})
    21  }
    22  
    23  func sendWithContextErrChecking(ctx context.Context, send func() error) error {
    24  	// If the context has been canceled or its deadline exceeded, we should return it
    25  	// instead of the cryptic error the Send() will return.
    26  	if ctxErr := ctx.Err(); ctxErr != nil {
    27  		return ctxErr
    28  	}
    29  
    30  	if err := send(); err != nil {
    31  		// Experimentally, we've seen the context switching to done after the Send()
    32  		// has been  called, so here we do recheck the context in case of error.
    33  		if ctxErr := ctx.Err(); ctxErr != nil {
    34  			return ctxErr
    35  		}
    36  
    37  		return err
    38  	}
    39  
    40  	return nil
    41  }