From d5f6dd1a46446ebb440811691a6ee8ce2443320d Mon Sep 17 00:00:00 2001
From: Simon Zolin <s.zolin@adguard.com>
Date: Wed, 15 May 2019 13:11:36 +0300
Subject: [PATCH] - dns query log: robust file flushing mechanism

Before this patch we could exit the process without waiting for
 file writing task to complete.
As a result a file could become corrupted or a large chunk of data
 could be missing.

Now the main thread either waits until file writing task completes
 or it writes log buffer to file itself.
---
 dnsforward/dnsforward.go    |  2 +-
 dnsforward/querylog.go      | 21 ++++++++++-----------
 dnsforward/querylog_file.go | 11 ++++++++++-
 3 files changed, 21 insertions(+), 13 deletions(-)

diff --git a/dnsforward/dnsforward.go b/dnsforward/dnsforward.go
index bd6cf75a..07f65084 100644
--- a/dnsforward/dnsforward.go
+++ b/dnsforward/dnsforward.go
@@ -236,7 +236,7 @@ func (s *Server) stopInternal() error {
 	}
 
 	// flush remainder to file
-	return s.queryLog.flushLogBuffer()
+	return s.queryLog.flushLogBuffer(true)
 }
 
 // IsRunning returns true if the DNS server is running
diff --git a/dnsforward/querylog.go b/dnsforward/querylog.go
index 14325354..7387e43a 100644
--- a/dnsforward/querylog.go
+++ b/dnsforward/querylog.go
@@ -30,6 +30,8 @@ type queryLog struct {
 
 	logBufferLock sync.RWMutex
 	logBuffer     []*logEntry
+	fileFlushLock sync.Mutex // synchronize a file-flushing goroutine and main thread
+	flushPending  bool       // don't start another goroutine while the previous one is still running
 
 	queryLogCache []*logEntry
 	queryLogLock  sync.RWMutex
@@ -91,13 +93,15 @@ func (l *queryLog) logRequest(question *dns.Msg, answer *dns.Msg, result *dnsfil
 		IP:       ip,
 		Upstream: upstream,
 	}
-	var flushBuffer []*logEntry
 
 	l.logBufferLock.Lock()
 	l.logBuffer = append(l.logBuffer, &entry)
-	if len(l.logBuffer) >= logBufferCap {
-		flushBuffer = l.logBuffer
-		l.logBuffer = nil
+	needFlush := false
+	if !l.flushPending {
+		needFlush = len(l.logBuffer) >= logBufferCap
+		if needFlush {
+			l.flushPending = true
+		}
 	}
 	l.logBufferLock.Unlock()
 	l.queryLogLock.Lock()
@@ -116,15 +120,10 @@ func (l *queryLog) logRequest(question *dns.Msg, answer *dns.Msg, result *dnsfil
 	}
 
 	// if buffer needs to be flushed to disk, do it now
-	if len(flushBuffer) > 0 {
+	if needFlush {
 		// write to file
 		// do it in separate goroutine -- we are stalling DNS response this whole time
-		go func() {
-			err := l.flushToFile(flushBuffer)
-			if err != nil {
-				log.Printf("Failed to flush the query log: %s", err)
-			}
-		}()
+		go l.flushLogBuffer(false)
 	}
 
 	return &entry
diff --git a/dnsforward/querylog_file.go b/dnsforward/querylog_file.go
index cb2ca40e..2147b042 100644
--- a/dnsforward/querylog_file.go
+++ b/dnsforward/querylog_file.go
@@ -20,11 +20,20 @@ var (
 const enableGzip = false
 
 // flushLogBuffer flushes the current buffer to file and resets the current buffer
-func (l *queryLog) flushLogBuffer() error {
+func (l *queryLog) flushLogBuffer(fullFlush bool) error {
+	l.fileFlushLock.Lock()
+	defer l.fileFlushLock.Unlock()
+
 	// flush remainder to file
 	l.logBufferLock.Lock()
+	needFlush := len(l.logBuffer) >= logBufferCap
+	if !needFlush && !fullFlush {
+		l.logBufferLock.Unlock()
+		return nil
+	}
 	flushBuffer := l.logBuffer
 	l.logBuffer = nil
+	l.flushPending = false
 	l.logBufferLock.Unlock()
 	err := l.flushToFile(flushBuffer)
 	if err != nil {