github.com/matrixorigin/matrixone@v1.2.0/pkg/fileservice/io_entry_reader.go (about)

     1  // Copyright 2022 Matrix Origin
     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 fileservice
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  
    21  	"github.com/matrixorigin/matrixone/pkg/common/moerr"
    22  )
    23  
    24  type ioEntriesReader struct {
    25  	ctx     context.Context
    26  	entries []IOEntry
    27  	offset  int64
    28  }
    29  
    30  var _ io.Reader = new(ioEntriesReader)
    31  
    32  func newIOEntriesReader(ctx context.Context, entries []IOEntry) *ioEntriesReader {
    33  	es := make([]IOEntry, len(entries))
    34  	copy(es, entries)
    35  	return &ioEntriesReader{
    36  		ctx:     ctx,
    37  		entries: es,
    38  	}
    39  }
    40  
    41  func (i *ioEntriesReader) Read(buf []byte) (n int, err error) {
    42  	for {
    43  
    44  		select {
    45  		case <-i.ctx.Done():
    46  			return n, i.ctx.Err()
    47  		default:
    48  		}
    49  
    50  		// no more data
    51  		if len(i.entries) == 0 {
    52  			err = io.EOF
    53  			return
    54  		}
    55  
    56  		entry := i.entries[0]
    57  
    58  		// gap
    59  		if i.offset < entry.Offset {
    60  			numBytes := entry.Offset - i.offset
    61  			if l := int64(len(buf)); l < numBytes {
    62  				numBytes = l
    63  			}
    64  			buf = buf[numBytes:] // skip
    65  			n += int(numBytes)
    66  			i.offset += numBytes
    67  		}
    68  
    69  		// buffer full
    70  		if len(buf) == 0 {
    71  			return
    72  		}
    73  
    74  		// copy data
    75  		var bytesRead int
    76  
    77  		if entry.ReaderForWrite != nil && entry.Size < 0 {
    78  			// read from size unknown reader
    79  			bytesRead, err = entry.ReaderForWrite.Read(buf)
    80  			i.entries[0].Offset += int64(bytesRead)
    81  			if err == io.EOF {
    82  				i.entries = i.entries[1:]
    83  			} else if err != nil {
    84  				return
    85  			}
    86  
    87  		} else if entry.ReaderForWrite != nil {
    88  			bytesToRead := entry.Size
    89  			if l := int64(len(buf)); bytesToRead > l {
    90  				bytesToRead = l
    91  			}
    92  			r := io.LimitReader(entry.ReaderForWrite, int64(bytesToRead))
    93  			bytesRead, err = io.ReadFull(r, buf[:bytesToRead])
    94  			if err != nil {
    95  				return
    96  			}
    97  			if int64(bytesRead) != bytesToRead {
    98  				err = moerr.NewSizeNotMatchNoCtx("")
    99  				return
   100  			}
   101  			i.entries[0].Offset += int64(bytesRead)
   102  			i.entries[0].Size -= int64(bytesRead)
   103  			if i.entries[0].Size == 0 {
   104  				i.entries = i.entries[1:]
   105  			}
   106  
   107  		} else {
   108  			bytesToRead := entry.Size
   109  			if l := int64(len(buf)); bytesToRead > l {
   110  				bytesToRead = l
   111  			}
   112  			if int64(len(entry.Data)) != entry.Size {
   113  				err = moerr.NewSizeNotMatchNoCtx("")
   114  				return
   115  			}
   116  			bytesRead = copy(buf, entry.Data[:bytesToRead])
   117  			if int64(bytesRead) != bytesToRead {
   118  				err = moerr.NewSizeNotMatchNoCtx("")
   119  				return
   120  			}
   121  			i.entries[0].Data = entry.Data[bytesRead:]
   122  			i.entries[0].Offset += int64(bytesRead)
   123  			i.entries[0].Size -= int64(bytesRead)
   124  			if i.entries[0].Size == 0 {
   125  				i.entries = i.entries[1:]
   126  			}
   127  		}
   128  
   129  		buf = buf[bytesRead:]
   130  		i.offset += int64(bytesRead)
   131  		n += int(bytesRead)
   132  
   133  	}
   134  }