Support for DB-specific loggers. (#659)

* Support for DB-specific loggers.

This commit adds support for overriding the global logger in specific
DBs. If no DB-specific logger is given, the global logger will be used
instead.

backup.go

@@ -64,7 +64,7 @@ func (stream *Stream) Backup(w io.Writer, since uint64) (uint64, error) {
 				var err error
 				valCopy, err = item.ValueCopy(nil)
 				if err != nil {
-					Errorf("Key [%x, %d]. Error while fetching value [%v]\n",
+					stream.db.opt.Errorf("Key [%x, %d]. Error while fetching value [%v]\n",
 						item.Key(), item.Version(), err)
 					return nil, err
 				}

db.go

@@ -150,7 +150,8 @@ func (db *DB) replayFunction() func(Entry, valuePointer) error {
 				lastCommit = txnTs
 			}
 			if lastCommit != txnTs {
-				Warningf("Found an incomplete txn at timestamp %d. Discarding it.\n", lastCommit)
+				db.opt.Warningf("Found an incomplete txn at timestamp %d. Discarding it.\n",
+					lastCommit)
 				txn = txn[:0]
 				lastCommit = txnTs
 			}
@@ -369,9 +370,9 @@ func (db *DB) Close() (err error) {
 	// We don't need to care about cstatus since no parallel compaction is running.
 	if db.opt.CompactL0OnClose {
 		if err := db.lc.doCompact(compactionPriority{level: 0, score: 1.73}); err != nil {
-			Warningf("Error while forcing compaction on level 0: %v", err)
+			db.opt.Warningf("Error while forcing compaction on level 0: %v", err)
 		} else {
-			Infof("Force compaction on level 0 done")
+			db.opt.Infof("Force compaction on level 0 done")
 		}
 	}
 
@@ -647,7 +648,7 @@ func (db *DB) doWrites(lc *y.Closer) {
 
 	writeRequests := func(reqs []*request) {
 		if err := db.writeRequests(reqs); err != nil {
-			Errorf("writeRequests: %v", err)
+			db.opt.Errorf("writeRequests: %v", err)
 		}
 		<-pendingCh
 	}
@@ -794,7 +795,7 @@ type flushTask struct {
 func (db *DB) handleFlushTask(ft flushTask) error {
 	if !ft.mt.Empty() {
 		// Store badger head even if vptr is zero, need it for readTs
-		Infof("Storing value log head: %+v\n", ft.vptr)
+		db.opt.Infof("Storing value log head: %+v\n", ft.vptr)
 		db.elog.Printf("Storing offset: %+v\n", ft.vptr)
 		offset := make([]byte, vptrSize)
 		ft.vptr.Encode(offset)
@@ -869,7 +870,7 @@ func (db *DB) flushMemtable(lc *y.Closer) error {
 				break
 			}
 			// Encounterd error. Retry indefinitely.
-			Errorf("Failure while flushing memtable to disk: %v. Retrying...\n", err)
+			db.opt.Errorf("Failure while flushing memtable to disk: %v. Retrying...\n", err)
 			time.Sleep(time.Second)
 		}
 	}
@@ -1166,7 +1167,7 @@ func (db *DB) Flatten(workers int) error {
 	defer db.startCompactions()
 
 	compactAway := func(cp compactionPriority) error {
-		Infof("Attempting to compact with %+v\n", cp)
+		db.opt.Infof("Attempting to compact with %+v\n", cp)
 		errCh := make(chan error, 1)
 		for i := 0; i < workers; i++ {
 			go func() {
@@ -1179,7 +1180,7 @@ func (db *DB) Flatten(workers int) error {
 			err := <-errCh
 			if err != nil {
 				rerr = err
-				Warningf("While running doCompact with %+v. Error: %v\n", cp, err)
+				db.opt.Warningf("While running doCompact with %+v. Error: %v\n", cp, err)
 			} else {
 				success++
 			}
@@ -1188,7 +1189,7 @@ func (db *DB) Flatten(workers int) error {
 			return rerr
 		}
 		// We could do at least one successful compaction. So, we'll consider this a success.
-		Infof("%d compactor(s) succeeded. One or more tables from level %d compacted.\n",
+		db.opt.Infof("%d compactor(s) succeeded. One or more tables from level %d compacted.\n",
 			success, cp.level)
 		return nil
 	}
@@ -1198,11 +1199,11 @@ func (db *DB) Flatten(workers int) error {
 	}
 
 	for {
-		Infof("\n")
+		db.opt.Infof("\n")
 		var levels []int
 		for i, l := range db.lc.levels {
 			sz := l.getTotalSize()
-			Infof("Level: %d. %8s Size. %8s Max.\n",
+			db.opt.Infof("Level: %d. %8s Size. %8s Max.\n",
 				i, hbytes(l.getTotalSize()), hbytes(l.maxTotalSize))
 			if sz > 0 {
 				levels = append(levels, i)
@@ -1211,7 +1212,7 @@ func (db *DB) Flatten(workers int) error {
 		if len(levels) == 1 {
 			prios := db.lc.pickCompactLevels()
 			if len(prios) == 0 || prios[0].score <= 1.0 {
-				Infof("All tables consolidated into one level. Flattening done.\n")
+				db.opt.Infof("All tables consolidated into one level. Flattening done.\n")
 				return nil
 			}
 			if err := compactAway(prios[0]); err != nil {
@@ -1243,18 +1244,18 @@ func (db *DB) DropAll() error {
 	if db.opt.ReadOnly {
 		panic("Attempting to drop data in read-only mode.")
 	}
-	Infof("DropAll called. Blocking writes...")
+	db.opt.Infof("DropAll called. Blocking writes...")
 	// Stop accepting new writes.
 	atomic.StoreInt32(&db.blockWrites, 1)
 
 	// Make all pending writes finish. The following will also close writeCh.
 	db.closers.writes.SignalAndWait()
-	Infof("Writes flushed. Stopping compactions now...")
+	db.opt.Infof("Writes flushed. Stopping compactions now...")
 
 	// Stop all compactions.
 	db.stopCompactions()
 	defer func() {
-		Infof("Resuming writes")
+		db.opt.Infof("Resuming writes")
 		db.startCompactions()
 
 		db.writeCh = make(chan *request, kvWriteChCapacity)
@@ -1264,7 +1265,7 @@ func (db *DB) DropAll() error {
 		// Resume writes.
 		atomic.StoreInt32(&db.blockWrites, 0)
 	}()
-	Infof("Compactions stopped. Dropping all SSTables...")
+	db.opt.Infof("Compactions stopped. Dropping all SSTables...")
 
 	// Remove inmemory tables. Calling DecrRef for safety. Not sure if they're absolutely needed.
 	db.mt.DecrRef()
@@ -1278,13 +1279,13 @@ func (db *DB) DropAll() error {
 	if err != nil {
 		return err
 	}
-	Infof("Deleted %d SSTables. Now deleting value logs...\n", num)
+	db.opt.Infof("Deleted %d SSTables. Now deleting value logs...\n", num)
 
 	num, err = db.vlog.dropAll()
 	if err != nil {
 		return err
 	}
 	db.vhead = valuePointer{} // Zero it out.
-	Infof("Deleted %d value log files. DropAll done.\n", num)
+	db.opt.Infof("Deleted %d value log files. DropAll done.\n", num)
 	return nil
 }

levels.go

@@ -241,7 +241,7 @@ func (s *levelsController) runWorker(lc *y.Closer) {
 				if err := s.doCompact(p); err == nil {
 					break
 				} else {
-					Errorf("Error while running doCompact: %v\n", err)
+					s.kv.opt.Errorf("Error while running doCompact: %v\n", err)
 				}
 			}
 		case <-lc.HasBeenClosed():

logger.go

@@ -30,26 +30,59 @@ type Logger interface {
 
 var badgerLogger Logger
 
+// SetLogger sets the global logger.
 func SetLogger(l Logger) { badgerLogger = l }
 
+// Errorf logs an ERROR message to the global logger.
 func Errorf(format string, v ...interface{}) {
 	badgerLogger.Errorf(format, v...)
 }
 
+// Errorf logs an ERROR log message to the logger specified in opts or to the
+// global logger if no logger is specified in opts.
+func (opt *Options) Errorf(format string, v ...interface{}) {
+	if opt.Logger != nil {
+		opt.Logger.Errorf(format, v...)
+		return
+	}
+	Errorf(format, v...)
+}
+
+// Infof logs an INFO message to the global logger.
 func Infof(format string, v ...interface{}) {
 	badgerLogger.Infof(format, v...)
 }
 
+// Infof is like Errorf but for INFO messages.
+func (opt *Options) Infof(format string, v ...interface{}) {
+	if opt.Logger != nil {
+		opt.Logger.Infof(format, v...)
+		return
+	}
+	Infof(format, v...)
+}
+
+// Warningf logs a WARNING message to the global logger.
 func Warningf(format string, v ...interface{}) {
 	badgerLogger.Warningf(format, v...)
 }
 
+// Warningf is like Errorf but for WARNING messages.
+func (opt *Options) Warningf(format string, v ...interface{}) {
+	if opt.Logger != nil {
+		opt.Logger.Warningf(format, v...)
+		return
+	}
+	Warningf(format, v...)
+}
+
 type defaultLog struct {
 	*log.Logger
 }
 
 var defaultLogger = &defaultLog{Logger: log.New(os.Stderr, "badger ", log.LstdFlags)}
 
+// UseDefaultLogger sets the global logger to the default logger.
 func UseDefaultLogger() { SetLogger(defaultLogger) }
 
 func (l *defaultLog) Errorf(f string, v ...interface{}) {

logger_test.go

@@ -0,0 +1,67 @@
+/*
+ * Copyright 2017 Dgraph Labs, Inc. and Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package badger
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type mockLogger struct {
+	output string
+}
+
+func (l *mockLogger) Errorf(f string, v ...interface{}) {
+	l.output = fmt.Sprintf("ERROR: "+f, v...)
+}
+
+func (l *mockLogger) Infof(f string, v ...interface{}) {
+	l.output = fmt.Sprintf("INFO: "+f, v...)
+}
+
+func (l *mockLogger) Warningf(f string, v ...interface{}) {
+	l.output = fmt.Sprintf("WARNING: "+f, v...)
+}
+
+// Test that the DB-specific log is used instead of the global log.
+func TestDbLog(t *testing.T) {
+	l := &mockLogger{}
+	opt := Options{Logger: l}
+
+	opt.Errorf("test")
+	require.Equal(t, "ERROR: test", l.output)
+	opt.Infof("test")
+	require.Equal(t, "INFO: test", l.output)
+	opt.Warningf("test")
+	require.Equal(t, "WARNING: test", l.output)
+}
+
+// Test that the global logger is used when no logger is specified in Options.
+func TestNoDbLog(t *testing.T) {
+	l := &mockLogger{}
+	SetLogger(l)
+	opt := Options{}
+
+	opt.Errorf("test")
+	require.Equal(t, "ERROR: test", l.output)
+	opt.Infof("test")
+	require.Equal(t, "INFO: test", l.output)
+	opt.Warningf("test")
+	require.Equal(t, "WARNING: test", l.output)
+}

merge.go

@@ -134,7 +134,7 @@ func (op *MergeOperator) runCompactions(dur time.Duration) {
 		case <-ticker.C: // wait for tick
 		}
 		if err := op.compact(); err != nil {
-			Errorf("failure while running merge operation: %s", err)
+			op.db.opt.Errorf("failure while running merge operation: %s", err)
 		}
 		if stop {
 			ticker.Stop()

options.go

@@ -106,6 +106,9 @@ type Options struct {
 
 	// Truncate value log to delete corrupt data, if any. Would not truncate if ReadOnly is set.
 	Truncate bool
+
+	// DB-specific logger which will override the global logger.
+	Logger Logger
 }
 
 // DefaultOptions sets a list of recommended options for good performance.
@@ -135,6 +138,7 @@ var DefaultOptions = Options{
 	ValueLogMaxEntries: 1000000,
 	ValueThreshold:     32,
 	Truncate:           false,
+	Logger:             defaultLogger,
 }
 
 // LSMOnlyOptions follows from DefaultOptions, but sets a higher ValueThreshold

stream.go

@@ -229,7 +229,7 @@ func (st *Stream) streamKVs(ctx context.Context) error {
 		if err := st.Send(batch); err != nil {
 			return err
 		}
-		Infof("%s Created batch of size: %s in %s.\n",
+		st.db.opt.Infof("%s Created batch of size: %s in %s.\n",
 			st.LogPrefix, humanize.Bytes(sz), time.Since(t))
 		return nil
 	}
@@ -248,7 +248,7 @@ outer:
 				continue
 			}
 			speed := bytesSent / durSec
-			Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec\n", st.LogPrefix,
+			st.db.opt.Infof("%s Time elapsed: %s, bytes sent: %s, speed: %s/sec\n", st.LogPrefix,
 				y.FixedDuration(dur), humanize.Bytes(bytesSent), humanize.Bytes(speed))
 
 		case kvs, ok := <-st.kvChan:
@@ -263,7 +263,7 @@ outer:
 		}
 	}
 
-	Infof("%s Sent %d keys\n", st.LogPrefix, count)
+	st.db.opt.Infof("%s Sent %d keys\n", st.LogPrefix, count)
 	return nil
 }
 

value.go

@@ -402,7 +402,7 @@ func (vlog *valueLog) rewrite(f *logFile, tr trace.Trace) error {
 				wb = wb[:0]
 			}
 		} else {
-			Warningf("This entry should have been caught. %+v\n", e)
+			vlog.db.opt.Warningf("This entry should have been caught. %+v\n", e)
 		}
 		return nil
 	}
@@ -420,7 +420,7 @@ func (vlog *valueLog) rewrite(f *logFile, tr trace.Trace) error {
 	for i := 0; i < len(wb); {
 		loops++
 		if batchSize == 0 {
-			Warningf("We shouldn't reach batch size of zero.")
+			vlog.db.opt.Warningf("We shouldn't reach batch size of zero.")
 			return ErrNoRewrite
 		}
 		end := i + batchSize
@@ -575,13 +575,13 @@ func (vlog *valueLog) dropAll() (int, error) {
 		loops++
 		if vlog.iteratorCount() == 0 {
 			if loops%10 == 0 {
-				Infof("Waiting for iterators to get done. Currently active: %d",
+				vlog.db.opt.Infof("Waiting for iterators to get done. Currently active: %d",
 					vlog.iteratorCount())
 			}
 			break
 		}
 	}
-	Infof("No active value log iterators. Deleting value logs...")
+	vlog.db.opt.Infof("No active value log iterators. Deleting value logs...")
 
 	var count int
 	deleteAll := func() error {
@@ -599,7 +599,7 @@ func (vlog *valueLog) dropAll() (int, error) {
 		return count, err
 	}
 
-	Infof("Value logs deleted. Creating value log file: 0")
+	vlog.db.opt.Infof("Value logs deleted. Creating value log file: 0")
 	if _, err := vlog.createVlogFile(0); err != nil {
 		return count, err
 	}
@@ -786,14 +786,14 @@ func (vlog *valueLog) open(db *DB, ptr valuePointer, replayFn logEntry) error {
 		if fid == ptr.Fid {
 			offset = ptr.Offset + ptr.Len
 		}
-		Infof("Replaying file id: %d at offset: %d\n", fid, offset)
+		vlog.db.opt.Infof("Replaying file id: %d at offset: %d\n", fid, offset)
 		now := time.Now()
 		// Replay and possible truncation done. Now we can open the file as per
 		// user specified options.
 		if err := vlog.replayLog(lf, offset, replayFn); err != nil {
 			return err
 		}
-		Infof("Replay took: %s\n", time.Since(now))
+		vlog.db.opt.Infof("Replay took: %s\n", time.Since(now))
 
 		if fid < vlog.maxFid {
 			if err := lf.openReadOnly(); err != nil {