github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/mmap_table_reader.go (about)

     1  // Copyright 2019 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  // This file incorporates work covered by the following copyright and
    16  // permission notice:
    17  //
    18  // Copyright 2016 Attic Labs, Inc. All rights reserved.
    19  // Licensed under the Apache License, version 2.0:
    20  // http://www.apache.org/licenses/LICENSE-2.0
    21  
    22  package nbs
    23  
    24  import (
    25  	"context"
    26  	"errors"
    27  	"fmt"
    28  	"io"
    29  	"math"
    30  	"os"
    31  	"path/filepath"
    32  	"strconv"
    33  	"time"
    34  
    35  	"github.com/dolthub/mmap-go"
    36  )
    37  
    38  type mmapTableReader struct {
    39  	tableReader
    40  	fc *fdCache
    41  	h  addr
    42  }
    43  
    44  const (
    45  	fileBlockSize = 1 << 12
    46  )
    47  
    48  var (
    49  	maxInt = int64(math.MaxInt64)
    50  )
    51  
    52  func init() {
    53  	if strconv.IntSize == 32 {
    54  		maxInt = math.MaxInt32
    55  	}
    56  }
    57  
    58  func newMmapTableReader(dir string, h addr, chunkCount uint32, indexCache *indexCache, fc *fdCache) (cs chunkSource, err error) {
    59  	path := filepath.Join(dir, h.String())
    60  
    61  	var index onHeapTableIndex
    62  	found := false
    63  	if indexCache != nil {
    64  		indexCache.lockEntry(h)
    65  		defer func() {
    66  			unlockErr := indexCache.unlockEntry(h)
    67  
    68  			if err == nil {
    69  				err = unlockErr
    70  			}
    71  		}()
    72  		index, found = indexCache.get(h)
    73  	}
    74  
    75  	if !found {
    76  		f := func() (ti onHeapTableIndex, err error) {
    77  			var f *os.File
    78  			f, err = fc.RefFile(path)
    79  
    80  			if err != nil {
    81  				return
    82  			}
    83  
    84  			defer func() {
    85  				unrefErr := fc.UnrefFile(path)
    86  
    87  				if unrefErr != nil {
    88  					err = unrefErr
    89  				}
    90  			}()
    91  
    92  			var fi os.FileInfo
    93  			fi, err = f.Stat()
    94  
    95  			if err != nil {
    96  				return
    97  			}
    98  
    99  			if fi.Size() < 0 {
   100  				// Size returns the number of bytes for regular files and is system dependant for others (Some of which can be negative).
   101  				err = fmt.Errorf("%s has invalid size: %d", path, fi.Size())
   102  				return
   103  			}
   104  
   105  			// index. Mmap won't take an offset that's not page-aligned, so find the nearest page boundary preceding the index.
   106  			indexOffset := fi.Size() - int64(footerSize) - int64(indexSize(chunkCount))
   107  			aligned := indexOffset / mmapAlignment * mmapAlignment // Thanks, integer arithmetic!
   108  
   109  			if fi.Size()-aligned > maxInt {
   110  				err = fmt.Errorf("%s - size: %d alignment: %d> maxInt: %d", path, fi.Size(), aligned, maxInt)
   111  				return
   112  			}
   113  
   114  			var mm mmap.MMap
   115  			mm, err = mmap.MapRegion(f, int(fi.Size()-aligned), mmap.RDONLY, 0, aligned)
   116  
   117  			if err != nil {
   118  				return
   119  			}
   120  
   121  			defer func() {
   122  				unmapErr := mm.Unmap()
   123  
   124  				if unmapErr != nil {
   125  					err = unmapErr
   126  				}
   127  			}()
   128  
   129  			buff := []byte(mm)
   130  			ti, err = parseTableIndex(buff[indexOffset-aligned:])
   131  
   132  			if err != nil {
   133  				return
   134  			}
   135  
   136  			if indexCache != nil {
   137  				indexCache.put(h, ti)
   138  			}
   139  
   140  			return
   141  		}
   142  
   143  		var err error
   144  		index, err = f()
   145  
   146  		if err != nil {
   147  			return nil, err
   148  		}
   149  	}
   150  
   151  	if chunkCount != index.chunkCount {
   152  		return nil, errors.New("unexpected chunk count")
   153  	}
   154  
   155  	return &mmapTableReader{
   156  		newTableReader(index, &cacheReaderAt{path, fc}, fileBlockSize),
   157  		fc,
   158  		h,
   159  	}, nil
   160  }
   161  
   162  func (mmtr *mmapTableReader) hash() (addr, error) {
   163  	return mmtr.h, nil
   164  }
   165  
   166  func (mmtr *mmapTableReader) Close() error {
   167  	return mmtr.tableReader.Close()
   168  }
   169  
   170  func (mmtr *mmapTableReader) Clone() chunkSource {
   171  	return &mmapTableReader{mmtr.tableReader.Clone(), mmtr.fc, mmtr.h}
   172  }
   173  
   174  type cacheReaderAt struct {
   175  	path string
   176  	fc   *fdCache
   177  }
   178  
   179  func (cra *cacheReaderAt) ReadAtWithStats(ctx context.Context, p []byte, off int64, stats *Stats) (n int, err error) {
   180  	var r io.ReaderAt
   181  	t1 := time.Now()
   182  
   183  	if r, err = cra.fc.RefFile(cra.path); err != nil {
   184  		return
   185  	}
   186  
   187  	defer func() {
   188  		stats.FileBytesPerRead.Sample(uint64(len(p)))
   189  		stats.FileReadLatency.SampleTimeSince(t1)
   190  	}()
   191  
   192  	defer func() {
   193  		unrefErr := cra.fc.UnrefFile(cra.path)
   194  
   195  		if err == nil {
   196  			err = unrefErr
   197  		}
   198  	}()
   199  
   200  	return r.ReadAt(p, off)
   201  }