github.com/coocood/badger@v1.5.1-0.20200528065104-c02ac3616d04/backup.go (about)

     1  /*
     2   * Copyright 2017 Dgraph Labs, Inc. and Contributors
     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 badger
    18  
    19  import (
    20  	"bufio"
    21  	"encoding/binary"
    22  	"io"
    23  	"log"
    24  	"sync"
    25  
    26  	"github.com/coocood/badger/y"
    27  
    28  	"github.com/coocood/badger/protos"
    29  )
    30  
    31  func writeTo(entry *protos.KVPair, w io.Writer) error {
    32  	if err := binary.Write(w, binary.LittleEndian, uint64(entry.Size())); err != nil {
    33  		return err
    34  	}
    35  	buf, err := entry.Marshal()
    36  	if err != nil {
    37  		return err
    38  	}
    39  	_, err = w.Write(buf)
    40  	return err
    41  }
    42  
    43  // Backup dumps a protobuf-encoded list of all entries in the database into the
    44  // given writer, that are newer than the specified version. It returns a
    45  // timestamp indicating when the entries were dumped which can be passed into a
    46  // later invocation to generate an incremental dump, of entries that have been
    47  // added/modified since the last invocation of DB.Backup()
    48  //
    49  // This can be used to backup the data in a database at a given point in time.
    50  func (db *DB) Backup(w io.Writer, since uint64) (uint64, error) {
    51  	var tsNew uint64
    52  	err := db.View(func(txn *Txn) error {
    53  		opts := DefaultIteratorOptions
    54  		opts.AllVersions = true
    55  		it := txn.NewIterator(opts)
    56  		defer it.Close()
    57  
    58  		for it.Rewind(); it.Valid(); it.Next() {
    59  			item := it.Item()
    60  			if item.Version() < since {
    61  				// Ignore versions less than given timestamp
    62  				continue
    63  			}
    64  			val, err := item.Value()
    65  			if err != nil {
    66  				log.Printf("Key [%x]. Error while fetching value [%v]\n", item.Key(), err)
    67  				continue
    68  			}
    69  
    70  			entry := &protos.KVPair{
    71  				Key:      y.Copy(item.Key()),
    72  				Value:    y.Copy(val),
    73  				UserMeta: item.UserMeta(),
    74  				Version:  item.Version(),
    75  			}
    76  
    77  			// Write entries to disk
    78  			if err := writeTo(entry, w); err != nil {
    79  				return err
    80  			}
    81  		}
    82  		tsNew = txn.readTs
    83  		return nil
    84  	})
    85  	return tsNew, err
    86  }
    87  
    88  // Load reads a protobuf-encoded list of all entries from a reader and writes
    89  // them to the database. This can be used to restore the database from a backup
    90  // made by calling DB.Backup().
    91  //
    92  // DB.Load() should be called on a database that is not running any other
    93  // concurrent transactions while it is running.
    94  func (db *DB) Load(r io.Reader) error {
    95  	br := bufio.NewReaderSize(r, 16<<10)
    96  	unmarshalBuf := make([]byte, 1<<10)
    97  	var entries []*Entry
    98  	var wg sync.WaitGroup
    99  	errChan := make(chan error, 1)
   100  
   101  	// func to check for pending error before sending off a batch for writing
   102  	batchSetAsyncIfNoErr := func(entries []*Entry) error {
   103  		select {
   104  		case err := <-errChan:
   105  			return err
   106  		default:
   107  			wg.Add(1)
   108  			return db.batchSetAsync(entries, func(err error) {
   109  				defer wg.Done()
   110  				if err != nil {
   111  					select {
   112  					case errChan <- err:
   113  					default:
   114  					}
   115  				}
   116  			})
   117  		}
   118  	}
   119  
   120  	for {
   121  		var sz uint64
   122  		err := binary.Read(br, binary.LittleEndian, &sz)
   123  		if err == io.EOF {
   124  			break
   125  		} else if err != nil {
   126  			return err
   127  		}
   128  
   129  		if cap(unmarshalBuf) < int(sz) {
   130  			unmarshalBuf = make([]byte, sz)
   131  		}
   132  
   133  		e := &protos.KVPair{}
   134  		if _, err = io.ReadFull(br, unmarshalBuf[:sz]); err != nil {
   135  			return err
   136  		}
   137  		if err = e.Unmarshal(unmarshalBuf[:sz]); err != nil {
   138  			return err
   139  		}
   140  		entries = append(entries, &Entry{
   141  			Key:      y.KeyWithTs(e.Key, e.Version),
   142  			Value:    e.Value,
   143  			UserMeta: e.UserMeta,
   144  		})
   145  		// Update nextCommit, memtable stores this timestamp in badger head
   146  		// when flushed.
   147  		if e.Version >= db.orc.commitTs() {
   148  			db.orc.nextCommit = e.Version + 1
   149  		}
   150  
   151  		if len(entries) == 1000 {
   152  			if err := batchSetAsyncIfNoErr(entries); err != nil {
   153  				return err
   154  			}
   155  			entries = make([]*Entry, 0, 1000)
   156  		}
   157  	}
   158  
   159  	if len(entries) > 0 {
   160  		if err := batchSetAsyncIfNoErr(entries); err != nil {
   161  			return err
   162  		}
   163  	}
   164  
   165  	wg.Wait()
   166  
   167  	select {
   168  	case err := <-errChan:
   169  		return err
   170  	default:
   171  		db.orc.curRead = db.orc.commitTs() - 1
   172  		return nil
   173  	}
   174  }