From 2d60f0c65a4a4e150751b1b3041068257e0254a6 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 27 Aug 2022 19:24:31 +0800
Subject: [PATCH] fix: logx disable not working in some cases (#2306)

* fix: logx disable not working in some cases

* fix: test fail
---
 core/logx/logs.go                    | 6 +++++-
 core/logx/logs_test.go               | 1 +
 core/stores/mon/collection_test.go   | 4 ----
 core/stores/mongo/collection_test.go | 4 ----
 4 files changed, 6 insertions(+), 9 deletions(-)

diff --git a/core/logx/logs.go b/core/logx/logs.go
index 451c8acc..76e75571 100644
--- a/core/logx/logs.go
+++ b/core/logx/logs.go
@@ -21,6 +21,7 @@ var (
 	logLevel   uint32
 	encoding   uint32 = jsonEncodingType
 	// use uint32 for atomic operations
+	disableLog  uint32
 	disableStat uint32
 	options     logOptions
 	writer      = new(atomicWriter)
@@ -73,6 +74,7 @@ func Close() error {
 
 // Disable disables the logging.
 func Disable() {
+	atomic.StoreUint32(&disableLog, 1)
 	writer.Store(nopWriter{})
 }
 
@@ -201,7 +203,9 @@ func SetLevel(level uint32) {
 
 // SetWriter sets the logging writer. It can be used to customize the logging.
 func SetWriter(w Writer) {
-	writer.Store(w)
+	if atomic.LoadUint32(&disableLog) == 0 {
+		writer.Store(w)
+	}
 }
 
 // SetUp sets up the logx. If already set up, just return nil.
diff --git a/core/logx/logs_test.go b/core/logx/logs_test.go
index 89faeb16..6c741539 100644
--- a/core/logx/logs_test.go
+++ b/core/logx/logs_test.go
@@ -599,6 +599,7 @@ func TestDisableStat(t *testing.T) {
 }
 
 func TestSetWriter(t *testing.T) {
+	atomic.StoreUint32(&disableLog, 0)
 	Reset()
 	SetWriter(nopWriter{})
 	assert.NotNil(t, writer.Load())
diff --git a/core/stores/mon/collection_test.go b/core/stores/mon/collection_test.go
index 75ddc493..225f008f 100644
--- a/core/stores/mon/collection_test.go
+++ b/core/stores/mon/collection_test.go
@@ -20,10 +20,6 @@ import (
 
 var errDummy = errors.New("dummy")
 
-func init() {
-	logx.Disable()
-}
-
 func TestKeepPromise_accept(t *testing.T) {
 	p := new(mockPromise)
 	kp := keepablePromise{
diff --git a/core/stores/mongo/collection_test.go b/core/stores/mongo/collection_test.go
index 4ae5f8c0..700fdcc2 100644
--- a/core/stores/mongo/collection_test.go
+++ b/core/stores/mongo/collection_test.go
@@ -17,10 +17,6 @@ import (
 
 var errDummy = errors.New("dummy")
 
-func init() {
-	logx.Disable()
-}
-
 func TestKeepPromise_accept(t *testing.T) {
 	p := new(mockPromise)
 	kp := keepablePromise{