github.com/hasnat/dolt/go@v0.0.0-20210628190320-9eb5d843fbb7/store/nbs/persisting_chunk_source.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  	"sync"
    30  	"time"
    31  
    32  	"golang.org/x/sync/errgroup"
    33  
    34  	"github.com/dolthub/dolt/go/store/atomicerr"
    35  	"github.com/dolthub/dolt/go/store/chunks"
    36  )
    37  
    38  var ErrNoReader = errors.New("could not get reader")
    39  var ErrNoChunkSource = errors.New("no chunk source")
    40  
    41  func newPersistingChunkSource(ctx context.Context, mt *memTable, haver chunkReader, p tablePersister, rl chan struct{}, stats *Stats) *persistingChunkSource {
    42  	t1 := time.Now()
    43  
    44  	ccs := &persistingChunkSource{ae: atomicerr.New(), mt: mt}
    45  	ccs.wg.Add(1)
    46  	rl <- struct{}{}
    47  	go func() {
    48  		defer ccs.wg.Done()
    49  		defer func() {
    50  			<-rl
    51  		}()
    52  
    53  		cs, err := p.Persist(ctx, mt, haver, stats)
    54  
    55  		if err != nil {
    56  			ccs.ae.SetIfError(err)
    57  			return
    58  		}
    59  
    60  		ccs.mu.Lock()
    61  		defer ccs.mu.Unlock()
    62  		ccs.cs = cs
    63  		ccs.mt = nil
    64  
    65  		cnt, err := cs.count()
    66  
    67  		if err != nil {
    68  			ccs.ae.SetIfError(err)
    69  			return
    70  		}
    71  
    72  		if cnt > 0 {
    73  			stats.PersistLatency.SampleTimeSince(t1)
    74  		}
    75  	}()
    76  
    77  	return ccs
    78  }
    79  
    80  type persistingChunkSource struct {
    81  	ae *atomicerr.AtomicError
    82  	mu sync.RWMutex
    83  	mt *memTable
    84  
    85  	wg sync.WaitGroup
    86  	cs chunkSource
    87  }
    88  
    89  func (ccs *persistingChunkSource) getReader() chunkReader {
    90  	ccs.mu.RLock()
    91  	defer ccs.mu.RUnlock()
    92  	if ccs.mt != nil {
    93  		return ccs.mt
    94  	}
    95  	return ccs.cs
    96  }
    97  
    98  func (ccs *persistingChunkSource) Close() error {
    99  	// persistingChunkSource does not own |cs| or |mt|. No need to close them.
   100  	return nil
   101  }
   102  
   103  func (ccs *persistingChunkSource) Clone() chunkSource {
   104  	// persistingChunkSource does not own |cs| or |mt|. No need to Clone.
   105  	return ccs
   106  }
   107  
   108  func (ccs *persistingChunkSource) has(h addr) (bool, error) {
   109  	cr := ccs.getReader()
   110  
   111  	if cr == nil {
   112  		return false, ErrNoReader
   113  	}
   114  
   115  	return cr.has(h)
   116  }
   117  
   118  func (ccs *persistingChunkSource) hasMany(addrs []hasRecord) (bool, error) {
   119  	cr := ccs.getReader()
   120  
   121  	if cr == nil {
   122  		return false, ErrNoReader
   123  	}
   124  	return cr.hasMany(addrs)
   125  }
   126  
   127  func (ccs *persistingChunkSource) get(ctx context.Context, h addr, stats *Stats) ([]byte, error) {
   128  	cr := ccs.getReader()
   129  
   130  	if cr == nil {
   131  		return nil, ErrNoReader
   132  	}
   133  
   134  	return cr.get(ctx, h, stats)
   135  }
   136  
   137  func (ccs *persistingChunkSource) getMany(ctx context.Context, eg *errgroup.Group, reqs []getRecord, found func(*chunks.Chunk), stats *Stats) (bool, error) {
   138  	cr := ccs.getReader()
   139  	if cr == nil {
   140  		return false, ErrNoReader
   141  	}
   142  	return cr.getMany(ctx, eg, reqs, found, stats)
   143  }
   144  
   145  func (ccs *persistingChunkSource) getManyCompressed(ctx context.Context, eg *errgroup.Group, reqs []getRecord, found func(CompressedChunk), stats *Stats) (bool, error) {
   146  	cr := ccs.getReader()
   147  	if cr == nil {
   148  		return false, ErrNoReader
   149  	}
   150  
   151  	return cr.getManyCompressed(ctx, eg, reqs, found, stats)
   152  }
   153  
   154  func (ccs *persistingChunkSource) wait() error {
   155  	ccs.wg.Wait()
   156  	return ccs.ae.Get()
   157  }
   158  
   159  func (ccs *persistingChunkSource) count() (uint32, error) {
   160  	err := ccs.wait()
   161  
   162  	if err != nil {
   163  		return 0, err
   164  	}
   165  
   166  	if ccs.cs == nil {
   167  		return 0, ErrNoChunkSource
   168  	}
   169  
   170  	return ccs.cs.count()
   171  }
   172  
   173  func (ccs *persistingChunkSource) uncompressedLen() (uint64, error) {
   174  	err := ccs.wait()
   175  
   176  	if err != nil {
   177  		return 0, err
   178  	}
   179  
   180  	if ccs.cs == nil {
   181  		return 0, ErrNoChunkSource
   182  	}
   183  
   184  	return ccs.cs.uncompressedLen()
   185  }
   186  
   187  func (ccs *persistingChunkSource) hash() (addr, error) {
   188  	err := ccs.wait()
   189  
   190  	if err != nil {
   191  		return addr{}, err
   192  	}
   193  
   194  	if ccs.cs == nil {
   195  		return addr{}, ErrNoChunkSource
   196  	}
   197  
   198  	return ccs.cs.hash()
   199  }
   200  
   201  func (ccs *persistingChunkSource) index() (tableIndex, error) {
   202  	err := ccs.wait()
   203  
   204  	if err != nil {
   205  		return onHeapTableIndex{}, err
   206  	}
   207  
   208  	if ccs.cs == nil {
   209  		return onHeapTableIndex{}, ErrNoChunkSource
   210  	}
   211  
   212  	return ccs.cs.index()
   213  }
   214  
   215  func (ccs *persistingChunkSource) reader(ctx context.Context) (io.Reader, error) {
   216  	err := ccs.wait()
   217  
   218  	if err != nil {
   219  		return nil, err
   220  	}
   221  
   222  	if ccs.cs == nil {
   223  		return nil, ErrNoChunkSource
   224  	}
   225  
   226  	return ccs.cs.reader(ctx)
   227  }
   228  
   229  func (ccs *persistingChunkSource) calcReads(reqs []getRecord, blockSize uint64) (reads int, remaining bool, err error) {
   230  	err = ccs.wait()
   231  
   232  	if err != nil {
   233  		return 0, false, err
   234  	}
   235  
   236  	if ccs.cs == nil {
   237  		return 0, false, ErrNoChunkSource
   238  	}
   239  
   240  	return ccs.cs.calcReads(reqs, blockSize)
   241  }
   242  
   243  func (ccs *persistingChunkSource) extract(ctx context.Context, chunks chan<- extractRecord) error {
   244  	err := ccs.wait()
   245  
   246  	if err != nil {
   247  		return err
   248  	}
   249  
   250  	if ccs.cs == nil {
   251  		return ErrNoChunkSource
   252  	}
   253  
   254  	return ccs.cs.extract(ctx, chunks)
   255  }
   256  
   257  type emptyChunkSource struct{}
   258  
   259  func (ecs emptyChunkSource) has(h addr) (bool, error) {
   260  	return false, nil
   261  }
   262  
   263  func (ecs emptyChunkSource) hasMany(addrs []hasRecord) (bool, error) {
   264  	return true, nil
   265  }
   266  
   267  func (ecs emptyChunkSource) get(ctx context.Context, h addr, stats *Stats) ([]byte, error) {
   268  	return nil, nil
   269  }
   270  
   271  func (ecs emptyChunkSource) getMany(ctx context.Context, eg *errgroup.Group, reqs []getRecord, found func(*chunks.Chunk), stats *Stats) (bool, error) {
   272  	return true, nil
   273  }
   274  
   275  func (ecs emptyChunkSource) getManyCompressed(ctx context.Context, eg *errgroup.Group, reqs []getRecord, found func(CompressedChunk), stats *Stats) (bool, error) {
   276  	return true, nil
   277  }
   278  
   279  func (ecs emptyChunkSource) count() (uint32, error) {
   280  	return 0, nil
   281  }
   282  
   283  func (ecs emptyChunkSource) uncompressedLen() (uint64, error) {
   284  	return 0, nil
   285  }
   286  
   287  func (ecs emptyChunkSource) hash() (addr, error) {
   288  	return addr{}, nil
   289  }
   290  
   291  func (ecs emptyChunkSource) index() (tableIndex, error) {
   292  	return onHeapTableIndex{}, nil
   293  }
   294  
   295  func (ecs emptyChunkSource) reader(context.Context) (io.Reader, error) {
   296  	return &bytes.Buffer{}, nil
   297  }
   298  
   299  func (ecs emptyChunkSource) calcReads(reqs []getRecord, blockSize uint64) (reads int, remaining bool, err error) {
   300  	return 0, true, nil
   301  }
   302  
   303  func (ecs emptyChunkSource) extract(ctx context.Context, chunks chan<- extractRecord) error {
   304  	return nil
   305  }
   306  
   307  func (ecs emptyChunkSource) Close() error {
   308  	return nil
   309  }
   310  
   311  func (ecs emptyChunkSource) Clone() chunkSource {
   312  	return ecs
   313  }