storj.io/minio@v0.0.0-20230509071714-0cbc90f649b1/pkg/s3select/simdj/reader.go (about)

     1  /*
     2   * MinIO Cloud Storage, (C) 2019 MinIO, Inc.
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package simdj
    18  
    19  import (
    20  	"fmt"
    21  	"io"
    22  	"sync"
    23  
    24  	"github.com/minio/simdjson-go"
    25  
    26  	"storj.io/minio/pkg/s3select/json"
    27  	"storj.io/minio/pkg/s3select/sql"
    28  )
    29  
    30  // Reader - JSON record reader for S3Select.
    31  type Reader struct {
    32  	args    *json.ReaderArgs
    33  	input   chan simdjson.Stream
    34  	decoded chan simdjson.Object
    35  
    36  	// err will only be returned after decoded has been closed.
    37  	err        *error
    38  	readCloser io.ReadCloser
    39  
    40  	exitReader chan struct{}
    41  	readerWg   sync.WaitGroup
    42  }
    43  
    44  // Read - reads single record.
    45  func (r *Reader) Read(dst sql.Record) (sql.Record, error) {
    46  	v, ok := <-r.decoded
    47  	if !ok {
    48  		if r.err != nil && *r.err != nil {
    49  			return nil, errJSONParsingError(*r.err)
    50  		}
    51  		return nil, io.EOF
    52  	}
    53  	dstRec, ok := dst.(*Record)
    54  	if !ok {
    55  		dstRec = &Record{}
    56  	}
    57  	dstRec.object = v
    58  	return dstRec, nil
    59  }
    60  
    61  // Close - closes underlying reader.
    62  func (r *Reader) Close() error {
    63  	// Close the input.
    64  	// Potentially racy if the stream decoder is still reading.
    65  	if r.readCloser != nil {
    66  		r.readCloser.Close()
    67  	}
    68  	if r.exitReader != nil {
    69  		close(r.exitReader)
    70  		r.readerWg.Wait()
    71  		r.exitReader = nil
    72  		r.input = nil
    73  	}
    74  	return nil
    75  }
    76  
    77  // startReader will start a reader that accepts input from r.input.
    78  // Input should be root -> object input. Each root indicates a record.
    79  // If r.input is closed, it is assumed that no more input will come.
    80  // When this function returns r.readerWg will be decremented and r.decoded will be closed.
    81  // On errors, r.err will be set. This should only be accessed after r.decoded has been closed.
    82  func (r *Reader) startReader() {
    83  	defer r.readerWg.Done()
    84  	defer close(r.decoded)
    85  	var tmpObj simdjson.Object
    86  	for {
    87  		var in simdjson.Stream
    88  		select {
    89  		case in = <-r.input:
    90  		case <-r.exitReader:
    91  			return
    92  		}
    93  		if in.Error != nil && in.Error != io.EOF {
    94  			r.err = &in.Error
    95  			return
    96  		}
    97  		if in.Value == nil {
    98  			if in.Error == io.EOF {
    99  				return
   100  			}
   101  			continue
   102  		}
   103  		i := in.Value.Iter()
   104  	readloop:
   105  		for {
   106  			var next simdjson.Iter
   107  			typ, err := i.AdvanceIter(&next)
   108  			if err != nil {
   109  				r.err = &err
   110  				return
   111  			}
   112  			switch typ {
   113  			case simdjson.TypeNone:
   114  				break readloop
   115  			case simdjson.TypeRoot:
   116  				typ, obj, err := next.Root(nil)
   117  				if err != nil {
   118  					r.err = &err
   119  					return
   120  				}
   121  				if typ != simdjson.TypeObject {
   122  					if typ == simdjson.TypeNone {
   123  						continue
   124  					}
   125  					err = fmt.Errorf("unexpected json type below root :%v", typ)
   126  					r.err = &err
   127  					return
   128  				}
   129  
   130  				o, err := obj.Object(&tmpObj)
   131  				if err != nil {
   132  					r.err = &err
   133  					return
   134  				}
   135  				select {
   136  				case <-r.exitReader:
   137  					return
   138  				case r.decoded <- *o:
   139  				}
   140  			default:
   141  				err = fmt.Errorf("unexpected root json type:%v", typ)
   142  				r.err = &err
   143  				return
   144  			}
   145  		}
   146  		if in.Error == io.EOF {
   147  			return
   148  		}
   149  	}
   150  }
   151  
   152  // NewReader - creates new JSON reader using readCloser.
   153  func NewReader(readCloser io.ReadCloser, args *json.ReaderArgs) *Reader {
   154  	r := Reader{
   155  		args:       args,
   156  		readCloser: readCloser,
   157  		decoded:    make(chan simdjson.Object, 1000),
   158  		input:      make(chan simdjson.Stream, 2),
   159  		exitReader: make(chan struct{}),
   160  	}
   161  	simdjson.ParseNDStream(readCloser, r.input, nil)
   162  	r.readerWg.Add(1)
   163  	go r.startReader()
   164  	return &r
   165  }
   166  
   167  // NewElementReader - creates new JSON reader using readCloser.
   168  func NewElementReader(ch chan simdjson.Object, err *error, args *json.ReaderArgs) *Reader {
   169  	return &Reader{
   170  		args:       args,
   171  		decoded:    ch,
   172  		err:        err,
   173  		readCloser: nil,
   174  	}
   175  }
   176  
   177  // NewTapeReaderChan will start a reader that will read input from the provided channel.
   178  func NewTapeReaderChan(pj chan simdjson.Stream, args *json.ReaderArgs) *Reader {
   179  	r := Reader{
   180  		args:       args,
   181  		decoded:    make(chan simdjson.Object, 1000),
   182  		input:      pj,
   183  		exitReader: make(chan struct{}),
   184  	}
   185  	r.readerWg.Add(1)
   186  	go r.startReader()
   187  	return &r
   188  }