From 115c6926b51ce6d0680b1232e2b48b40a9284f73 Mon Sep 17 00:00:00 2001
From: Sachin Maurya <57769917+slayer321@users.noreply.github.com>
Date: Mon, 27 Nov 2023 11:45:43 +0530
Subject: [PATCH] fix(badger): Add missing SamplingStoreFactory.CreateLock
 method (#4966)

---
 plugin/storage/badger/factory.go      | 10 +++++---
 plugin/storage/badger/factory_test.go |  4 +++
 plugin/storage/badger/lock.go         | 29 +++++++++++++++++++++
 plugin/storage/badger/lock_test.go    | 36 +++++++++++++++++++++++++++
 plugin/storage/factory_test.go        |  2 +-
 5 files changed, 77 insertions(+), 4 deletions(-)
 create mode 100644 plugin/storage/badger/lock.go
 create mode 100644 plugin/storage/badger/lock_test.go

diff --git a/plugin/storage/badger/factory.go b/plugin/storage/badger/factory.go
index 40f90329c5b..ee0ad29c2e0 100644
--- a/plugin/storage/badger/factory.go
+++ b/plugin/storage/badger/factory.go
@@ -26,6 +26,7 @@ import (
 	"github.com/spf13/viper"
 	"go.uber.org/zap"
 
+	"github.com/jaegertracing/jaeger/pkg/distributedlock"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
 	"github.com/jaegertracing/jaeger/plugin"
 	depStore "github.com/jaegertracing/jaeger/plugin/storage/badger/dependencystore"
@@ -52,9 +53,7 @@ var ( // interface comformance checks
 	// TODO badger could implement archive storage
 	// _ storage.ArchiveFactory       = (*Factory)(nil)
 
-	// TODO CreateLock function is missing
-	// Being fixed in https://github.com/jaegertracing/jaeger/pull/4966
-	// _ storage.SamplingStoreFactory = (*Factory)(nil)
+	_ storage.SamplingStoreFactory = (*Factory)(nil)
 )
 
 // Factory implements storage.Factory for Badger backend.
@@ -186,6 +185,11 @@ func (f *Factory) CreateSamplingStore(maxBuckets int) (samplingstore.Store, erro
 	return badgerSampling.NewSamplingStore(f.store), nil
 }
 
+// CreateLock implements storage.SamplingStoreFactory
+func (f *Factory) CreateLock() (distributedlock.Lock, error) {
+	return &lock{}, nil
+}
+
 // Close Implements io.Closer and closes the underlying storage
 func (f *Factory) Close() error {
 	close(f.maintenanceDone)
diff --git a/plugin/storage/badger/factory_test.go b/plugin/storage/badger/factory_test.go
index 9a78c0bb4f1..bf995e0666d 100644
--- a/plugin/storage/badger/factory_test.go
+++ b/plugin/storage/badger/factory_test.go
@@ -68,6 +68,10 @@ func TestForCodecov(t *testing.T) {
 	_, err = f.CreateDependencyReader()
 	assert.NoError(t, err)
 
+	lock, err := f.CreateLock()
+	assert.NoError(t, err)
+	assert.NotNil(t, lock)
+
 	// Now, remove the badger directories
 	err = os.RemoveAll(f.tmpDir)
 	assert.NoError(t, err)
diff --git a/plugin/storage/badger/lock.go b/plugin/storage/badger/lock.go
new file mode 100644
index 00000000000..036448f1965
--- /dev/null
+++ b/plugin/storage/badger/lock.go
@@ -0,0 +1,29 @@
+// Copyright (c) 2023 The Jaeger Authors.
+//
+// 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 "time"
+
+type lock struct{}
+
+// Acquire always returns true for badgerdb as no lock is needed
+func (l *lock) Acquire(resource string, ttl time.Duration) (bool, error) {
+	return true, nil
+}
+
+// Forfeit always returns true for badgerdb as no lock is needed
+func (l *lock) Forfeit(resource string) (bool, error) {
+	return true, nil
+}
diff --git a/plugin/storage/badger/lock_test.go b/plugin/storage/badger/lock_test.go
new file mode 100644
index 00000000000..f973104f9e6
--- /dev/null
+++ b/plugin/storage/badger/lock_test.go
@@ -0,0 +1,36 @@
+// Copyright (c) 2023 The Jaeger Authors.
+//
+// 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 (
+	"testing"
+	"time"
+
+	"github.com/crossdock/crossdock-go/assert"
+)
+
+func TestAcquire(t *testing.T) {
+	l := &lock{}
+	ok, err := l.Acquire("resource", time.Duration(1))
+	assert.True(t, ok)
+	assert.NoError(t, err)
+}
+
+func TestForfeit(t *testing.T) {
+	l := &lock{}
+	ok, err := l.Forfeit("resource")
+	assert.True(t, ok)
+	assert.NoError(t, err)
+}
diff --git a/plugin/storage/factory_test.go b/plugin/storage/factory_test.go
index 8aad9e35334..93390abf722 100644
--- a/plugin/storage/factory_test.go
+++ b/plugin/storage/factory_test.go
@@ -294,7 +294,7 @@ func TestCreateError(t *testing.T) {
 }
 
 func TestAllSamplingStorageTypes(t *testing.T) {
-	assert.Equal(t, AllSamplingStorageTypes(), []string{"cassandra", "memory"})
+	assert.Equal(t, []string{"cassandra", "memory", "badger"}, AllSamplingStorageTypes())
 }
 
 func TestCreateSamplingStoreFactory(t *testing.T) {