github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/file_table_persister.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  	"bytes"
    26  	"context"
    27  	"errors"
    28  	"io"
    29  	"io/ioutil"
    30  	"os"
    31  	"path"
    32  	"path/filepath"
    33  	"strings"
    34  
    35  	"github.com/dolthub/dolt/go/store/util/tempfiles"
    36  
    37  	"github.com/dolthub/dolt/go/store/d"
    38  )
    39  
    40  const tempTablePrefix = "nbs_table_"
    41  
    42  func newFSTablePersister(dir string, fc *fdCache, indexCache *indexCache) tablePersister {
    43  	d.PanicIfTrue(fc == nil)
    44  	return &fsTablePersister{dir, fc, indexCache}
    45  }
    46  
    47  type fsTablePersister struct {
    48  	dir        string
    49  	fc         *fdCache
    50  	indexCache *indexCache
    51  }
    52  
    53  func (ftp *fsTablePersister) Open(ctx context.Context, name addr, chunkCount uint32, stats *Stats) (chunkSource, error) {
    54  	return newMmapTableReader(ftp.dir, name, chunkCount, ftp.indexCache, ftp.fc)
    55  }
    56  
    57  func (ftp *fsTablePersister) Persist(ctx context.Context, mt *memTable, haver chunkReader, stats *Stats) (chunkSource, error) {
    58  	name, data, chunkCount, err := mt.write(haver, stats)
    59  
    60  	if err != nil {
    61  		return emptyChunkSource{}, err
    62  	}
    63  
    64  	return ftp.persistTable(ctx, name, data, chunkCount, stats)
    65  }
    66  
    67  func (ftp *fsTablePersister) persistTable(ctx context.Context, name addr, data []byte, chunkCount uint32, stats *Stats) (cs chunkSource, err error) {
    68  	if chunkCount == 0 {
    69  		return emptyChunkSource{}, nil
    70  	}
    71  
    72  	tempName, err := func() (tempName string, ferr error) {
    73  		var temp *os.File
    74  		temp, ferr = tempfiles.MovableTempFileProvider.NewFile(ftp.dir, tempTablePrefix)
    75  
    76  		if ferr != nil {
    77  			return "", ferr
    78  		}
    79  
    80  		defer func() {
    81  			closeErr := temp.Close()
    82  
    83  			if ferr == nil {
    84  				ferr = closeErr
    85  			}
    86  		}()
    87  
    88  		_, ferr = io.Copy(temp, bytes.NewReader(data))
    89  
    90  		if ferr != nil {
    91  			return "", ferr
    92  		}
    93  
    94  		index, ferr := parseTableIndex(data)
    95  
    96  		if ferr != nil {
    97  			return "", ferr
    98  		}
    99  
   100  		if ftp.indexCache != nil {
   101  			ftp.indexCache.lockEntry(name)
   102  			defer func() {
   103  				unlockErr := ftp.indexCache.unlockEntry(name)
   104  
   105  				if ferr == nil {
   106  					ferr = unlockErr
   107  				}
   108  			}()
   109  			ftp.indexCache.put(name, index)
   110  		}
   111  
   112  		return temp.Name(), nil
   113  	}()
   114  
   115  	if err != nil {
   116  		return nil, err
   117  	}
   118  
   119  	newName := filepath.Join(ftp.dir, name.String())
   120  	err = ftp.fc.ShrinkCache()
   121  
   122  	if err != nil {
   123  		return nil, err
   124  	}
   125  
   126  	err = os.Rename(tempName, newName)
   127  
   128  	if err != nil {
   129  		return nil, err
   130  	}
   131  
   132  	return ftp.Open(ctx, name, chunkCount, stats)
   133  }
   134  
   135  func (ftp *fsTablePersister) ConjoinAll(ctx context.Context, sources chunkSources, stats *Stats) (chunkSource, error) {
   136  	plan, err := planConjoin(sources, stats)
   137  
   138  	if err != nil {
   139  		return emptyChunkSource{}, err
   140  	}
   141  
   142  	if plan.chunkCount == 0 {
   143  		return emptyChunkSource{}, nil
   144  	}
   145  
   146  	name := nameFromSuffixes(plan.suffixes())
   147  	tempName, err := func() (tempName string, ferr error) {
   148  		var temp *os.File
   149  		temp, ferr = tempfiles.MovableTempFileProvider.NewFile(ftp.dir, tempTablePrefix)
   150  
   151  		if ferr != nil {
   152  			return "", ferr
   153  		}
   154  
   155  		defer func() {
   156  			closeErr := temp.Close()
   157  
   158  			if ferr == nil {
   159  				ferr = closeErr
   160  			}
   161  		}()
   162  
   163  		for _, sws := range plan.sources.sws {
   164  			var r io.Reader
   165  			r, ferr = sws.source.reader(ctx)
   166  
   167  			if ferr != nil {
   168  				return "", ferr
   169  			}
   170  
   171  			n, ferr := io.CopyN(temp, r, int64(sws.dataLen))
   172  
   173  			if ferr != nil {
   174  				return "", ferr
   175  			}
   176  
   177  			if uint64(n) != sws.dataLen {
   178  				return "", errors.New("failed to copy all data")
   179  			}
   180  		}
   181  
   182  		_, ferr = temp.Write(plan.mergedIndex)
   183  
   184  		if ferr != nil {
   185  			return "", ferr
   186  		}
   187  
   188  		var index onHeapTableIndex
   189  		index, ferr = parseTableIndex(plan.mergedIndex)
   190  
   191  		if ferr != nil {
   192  			return "", ferr
   193  		}
   194  
   195  		if ftp.indexCache != nil {
   196  			ftp.indexCache.put(name, index)
   197  		}
   198  
   199  		return temp.Name(), nil
   200  	}()
   201  
   202  	if err != nil {
   203  		return nil, err
   204  	}
   205  
   206  	err = os.Rename(tempName, filepath.Join(ftp.dir, name.String()))
   207  
   208  	if err != nil {
   209  		return nil, err
   210  	}
   211  
   212  	return ftp.Open(ctx, name, plan.chunkCount, stats)
   213  }
   214  
   215  func (ftp *fsTablePersister) PruneTableFiles(ctx context.Context, contents manifestContents) error {
   216  	ss := contents.getSpecSet()
   217  
   218  	fileInfos, err := ioutil.ReadDir(ftp.dir)
   219  
   220  	if err != nil {
   221  		return err
   222  	}
   223  
   224  	err = ftp.fc.ShrinkCache()
   225  
   226  	if err != nil {
   227  		return err
   228  	}
   229  
   230  	ea := make(gcErrAccum)
   231  	for _, info := range fileInfos {
   232  		if info.IsDir() {
   233  			continue
   234  		}
   235  
   236  		filePath := path.Join(ftp.dir, info.Name())
   237  
   238  		if strings.HasPrefix(info.Name(), tempTablePrefix) {
   239  			err = os.Remove(filePath)
   240  			if err != nil {
   241  				ea.add(filePath, err)
   242  			}
   243  			continue
   244  		}
   245  
   246  		if len(info.Name()) != 32 {
   247  			continue // not a table file
   248  		}
   249  
   250  		addy, err := parseAddr(info.Name())
   251  		if err != nil {
   252  			continue // not a table file
   253  		}
   254  
   255  		if _, ok := ss[addy]; ok {
   256  			continue // file is referenced in the manifest
   257  		}
   258  
   259  		err = os.Remove(filePath)
   260  		if err != nil {
   261  			ea.add(filePath, err)
   262  		}
   263  	}
   264  
   265  	if !ea.isEmpty() {
   266  		return ea
   267  	}
   268  
   269  	return nil
   270  }