github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/posting/writer.go (about)

     1  /*
     2   * Copyright 2018 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 posting
    18  
    19  import (
    20  	"math"
    21  	"sync"
    22  
    23  	"github.com/dgraph-io/badger"
    24  	"github.com/dgraph-io/badger/pb"
    25  	"github.com/golang/glog"
    26  )
    27  
    28  // TxnWriter is in charge or writing transactions to badger.
    29  type TxnWriter struct {
    30  	db  *badger.DB
    31  	wg  sync.WaitGroup
    32  	che chan error
    33  }
    34  
    35  // NewTxnWriter returns a new TxnWriter instance.
    36  func NewTxnWriter(db *badger.DB) *TxnWriter {
    37  	return &TxnWriter{
    38  		db:  db,
    39  		che: make(chan error, 1),
    40  	}
    41  }
    42  
    43  func (w *TxnWriter) cb(err error) {
    44  	defer w.wg.Done()
    45  	if err == nil {
    46  		return
    47  	}
    48  
    49  	glog.Errorf("TxnWriter got error during callback: %v", err)
    50  	select {
    51  	case w.che <- err:
    52  	default:
    53  	}
    54  }
    55  
    56  // Write stores the given key-value pairs in badger.
    57  func (w *TxnWriter) Write(kvs *pb.KVList) error {
    58  	for _, kv := range kvs.Kv {
    59  		var meta byte
    60  		if len(kv.UserMeta) > 0 {
    61  			meta = kv.UserMeta[0]
    62  		}
    63  		if err := w.SetAt(kv.Key, kv.Value, meta, kv.Version); err != nil {
    64  			return err
    65  		}
    66  	}
    67  	return nil
    68  }
    69  
    70  func (w *TxnWriter) update(commitTs uint64, f func(txn *badger.Txn) error) error {
    71  	if commitTs == 0 {
    72  		return nil
    73  	}
    74  	txn := w.db.NewTransactionAt(math.MaxUint64, true)
    75  	defer txn.Discard()
    76  
    77  	err := f(txn)
    78  	if err == badger.ErrTxnTooBig {
    79  		// continue to commit.
    80  	} else if err != nil {
    81  		return err
    82  	}
    83  	w.wg.Add(1)
    84  	return txn.CommitAt(commitTs, w.cb)
    85  }
    86  
    87  // SetAt writes a key-value pair at the given timestamp.
    88  func (w *TxnWriter) SetAt(key, val []byte, meta byte, ts uint64) error {
    89  	return w.update(ts, func(txn *badger.Txn) error {
    90  		switch meta {
    91  		case BitCompletePosting, BitEmptyPosting:
    92  			err := txn.SetEntry((&badger.Entry{
    93  				Key:      key,
    94  				Value:    val,
    95  				UserMeta: meta,
    96  			}).WithDiscard())
    97  			if err != nil {
    98  				return err
    99  			}
   100  		default:
   101  			err := txn.SetEntry(&badger.Entry{
   102  				Key:      key,
   103  				Value:    val,
   104  				UserMeta: meta,
   105  			})
   106  			if err != nil {
   107  				return err
   108  			}
   109  		}
   110  		return nil
   111  	})
   112  }
   113  
   114  // Flush waits until all operations are done and all data is written to disk.
   115  func (w *TxnWriter) Flush() error {
   116  	defer func() {
   117  		if err := w.db.Sync(); err != nil {
   118  			glog.Errorf("Error while calling Sync from TxnWriter.Flush: %v", err)
   119  		}
   120  	}()
   121  	w.wg.Wait()
   122  	select {
   123  	case err := <-w.che:
   124  		return err
   125  	default:
   126  		return nil
   127  	}
   128  }