github.com/minio/minio@v0.0.0-20240328213742-3f72439b8a27/internal/grid/stream.go (about)

     1  // Copyright (c) 2015-2023 MinIO, Inc.
     2  //
     3  // This file is part of MinIO Object Storage stack
     4  //
     5  // This program is free software: you can redistribute it and/or modify
     6  // it under the terms of the GNU Affero General Public License as published by
     7  // the Free Software Foundation, either version 3 of the License, or
     8  // (at your option) any later version.
     9  //
    10  // This program is distributed in the hope that it will be useful
    11  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    12  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13  // GNU Affero General Public License for more details.
    14  //
    15  // You should have received a copy of the GNU Affero General Public License
    16  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17  
    18  package grid
    19  
    20  import (
    21  	"context"
    22  	"errors"
    23  )
    24  
    25  // A Stream is a two-way stream.
    26  // All responses *must* be read by the caller.
    27  // If the call is canceled through the context,
    28  // the appropriate error will be returned.
    29  type Stream struct {
    30  	// responses from the remote server.
    31  	// Channel will be closed after error or when remote closes.
    32  	// All responses *must* be read by the caller until either an error is returned or the channel is closed.
    33  	// Canceling the context will cause the context cancellation error to be returned.
    34  	responses <-chan Response
    35  	cancel    context.CancelCauseFunc
    36  
    37  	// Requests sent to the server.
    38  	// If the handler is defined with 0 incoming capacity this will be nil.
    39  	// Channel *must* be closed to signal the end of the stream.
    40  	// If the request context is canceled, the stream will no longer process requests.
    41  	// Requests sent cannot be used any further by the called.
    42  	Requests chan<- []byte
    43  
    44  	muxID uint64
    45  	ctx   context.Context
    46  }
    47  
    48  // Send a payload to the remote server.
    49  func (s *Stream) Send(b []byte) error {
    50  	if s.Requests == nil {
    51  		return errors.New("stream does not accept requests")
    52  	}
    53  	select {
    54  	case s.Requests <- b:
    55  		return nil
    56  	case <-s.ctx.Done():
    57  		return context.Cause(s.ctx)
    58  	}
    59  }
    60  
    61  // Results returns the results from the remote server one by one.
    62  // If any error is returned by the callback, the stream will be canceled.
    63  // If the context is canceled, the stream will be canceled.
    64  func (s *Stream) Results(next func(b []byte) error) (err error) {
    65  	done := false
    66  	defer func() {
    67  		if s.cancel != nil {
    68  			s.cancel(err)
    69  		}
    70  
    71  		if !done {
    72  			// Drain channel.
    73  			for range s.responses {
    74  			}
    75  		}
    76  	}()
    77  	doneCh := s.ctx.Done()
    78  	for {
    79  		select {
    80  		case <-doneCh:
    81  			if err := context.Cause(s.ctx); !errors.Is(err, errStreamEOF) {
    82  				return err
    83  			}
    84  			// Fall through to be sure we have returned all responses.
    85  			doneCh = nil
    86  		case resp, ok := <-s.responses:
    87  			if !ok {
    88  				done = true
    89  				return nil
    90  			}
    91  			if resp.Err != nil {
    92  				return resp.Err
    93  			}
    94  			err = next(resp.Msg)
    95  			if err != nil {
    96  				return err
    97  			}
    98  		}
    99  	}
   100  }