github.com/waldiirawan/apm-agent-go/v2@v2.2.2/internal/iochan/reader.go (about)

     1  // Licensed to Elasticsearch B.V. under one or more contributor
     2  // license agreements. See the NOTICE file distributed with
     3  // this work for additional information regarding copyright
     4  // ownership. Elasticsearch B.V. licenses this file to you under
     5  // the Apache License, Version 2.0 (the "License"); you may
     6  // not use this file except in compliance with the License.
     7  // You may obtain a copy of the License at
     8  //
     9  //     http://www.apache.org/licenses/LICENSE-2.0
    10  //
    11  // Unless required by applicable law or agreed to in writing,
    12  // software distributed under the License is distributed on an
    13  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    14  // KIND, either express or implied.  See the License for the
    15  // specific language governing permissions and limitations
    16  // under the License.
    17  
    18  package iochan
    19  
    20  import (
    21  	"sync"
    22  )
    23  
    24  // Reader is a channel-based io.Reader.
    25  //
    26  // Reader is safe for use in a single producer, single consumer pattern.
    27  type Reader struct {
    28  	// C can be used for receiving read requests.
    29  	//
    30  	// Once a read request is received, it must be responded
    31  	// to, in order to avoid blocking the reader.
    32  	C    <-chan ReadRequest
    33  	c    chan ReadRequest
    34  	resp chan readResponse
    35  
    36  	mu          sync.Mutex
    37  	readClosed  bool
    38  	writeClosed bool
    39  	readErr     error
    40  }
    41  
    42  // NewReader returns a new Reader.
    43  func NewReader() *Reader {
    44  	c := make(chan ReadRequest)
    45  	return &Reader{
    46  		C:    c,
    47  		c:    c,
    48  		resp: make(chan readResponse, 1),
    49  	}
    50  }
    51  
    52  // CloseWrite closes reader.C. CloseWrite is idempotent,
    53  // but must not be called concurrently with Read.
    54  func (r *Reader) CloseWrite() {
    55  	r.mu.Lock()
    56  	defer r.mu.Unlock()
    57  	if !r.writeClosed {
    58  		r.writeClosed = true
    59  		close(r.c)
    60  	}
    61  }
    62  
    63  // CloseRead closes the reader such that any waiting or future
    64  // Reads return err. Additional calls to CloseRead have no
    65  // effect. CloseRead must not be called concurrently with
    66  // ReadRequest.Respond.
    67  func (r *Reader) CloseRead(err error) error {
    68  	r.mu.Lock()
    69  	defer r.mu.Unlock()
    70  	if !r.readClosed {
    71  		r.readClosed = true
    72  		r.readErr = err
    73  		close(r.resp)
    74  	}
    75  	return nil
    76  }
    77  
    78  // Read sends a ReadRequest to r.C containing buf, and returns the
    79  // response sent by the channel consumer via the read request's
    80  // Response method.
    81  func (r *Reader) Read(buf []byte) (int, error) {
    82  	select {
    83  	case <-r.resp:
    84  		return 0, r.readErr
    85  	case r.c <- ReadRequest{Buf: buf, response: r.resp}:
    86  	}
    87  	resp, ok := <-r.resp
    88  	if !ok {
    89  		return 0, r.readErr
    90  	}
    91  	return resp.N, resp.Err
    92  }
    93  
    94  // ReadRequest holds the buffer and response channel for a read request.
    95  type ReadRequest struct {
    96  	// Buf is the read buffer into which data should be read.
    97  	Buf      []byte
    98  	response chan<- readResponse
    99  }
   100  
   101  // Respond responds to the Read request. Respond must not be called
   102  // concurrently with Reader.Close.
   103  func (rr *ReadRequest) Respond(n int, err error) {
   104  	rr.response <- readResponse{N: n, Err: err}
   105  }
   106  
   107  type readResponse struct {
   108  	N   int
   109  	Err error
   110  }