github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/libraries/utils/async/async_reader.go (about)

     1  // Copyright 2020 Dolthub, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package async
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  	"sync"
    21  )
    22  
    23  // ReadFunc is a function that is called repeatedly in order to retrieve a stream of objects.  When all objects have been
    24  // been read from the stream then (nil, io.EOF) should be returned.
    25  type ReadFunc func(ctx context.Context) (interface{}, error)
    26  
    27  type objErrTuple struct {
    28  	obj interface{}
    29  	err error
    30  }
    31  
    32  // AsyncReader is a TableReadCloser implementation that spins up a go routine to keep reading data into
    33  // a buffer so that it is ready when the caller wants it.
    34  type AsyncReader struct {
    35  	readFunc ReadFunc
    36  	stopCh   chan struct{}
    37  	rowCh    chan objErrTuple
    38  	wg       *sync.WaitGroup
    39  }
    40  
    41  // NewAsyncReader creates a new AsyncReader
    42  func NewAsyncReader(rf ReadFunc, bufferSize int) *AsyncReader {
    43  	return &AsyncReader{rf, make(chan struct{}), make(chan objErrTuple, bufferSize), &sync.WaitGroup{}}
    44  }
    45  
    46  // Start the worker routine reading rows to the channel
    47  func (asRd *AsyncReader) Start(ctx context.Context) error {
    48  	asRd.wg.Add(1)
    49  	go func() {
    50  		defer asRd.wg.Done()
    51  		defer close(asRd.rowCh)
    52  		asRd.readObjects(ctx)
    53  	}()
    54  
    55  	return nil
    56  }
    57  
    58  // ReadObject reads an object
    59  func (asRd *AsyncReader) Read() (interface{}, error) {
    60  	objErrTup := <-asRd.rowCh
    61  
    62  	if objErrTup.obj == nil && objErrTup.err == nil {
    63  		return nil, io.EOF
    64  	}
    65  
    66  	return objErrTup.obj, objErrTup.err
    67  }
    68  
    69  // Close releases resources being held
    70  func (asRd *AsyncReader) Close() error {
    71  	close(asRd.stopCh)
    72  	asRd.wg.Wait()
    73  
    74  	return nil
    75  }
    76  
    77  // background read loop running in separate go routine
    78  func (asRd *AsyncReader) readObjects(ctx context.Context) {
    79  	for {
    80  		select {
    81  		case <-asRd.stopCh:
    82  			return
    83  		default:
    84  		}
    85  
    86  		obj, err := asRd.readFunc(ctx)
    87  		asRd.rowCh <- objErrTuple{obj, err}
    88  
    89  		if err != nil {
    90  			break
    91  		}
    92  	}
    93  }