[prober] Minor refactoring. (cloudprober#625)

- Use read-write mutex, instead of regular mutex, to control access to cloudprober config info. - Get raw and parsed config from the config source directly instead of storing it ahead of time. - Allow access to the prober instance.
jumpojoy · Nov 6, 2023 · 1e70553 · 1e70553
1 parent 20dac79
commit 1e70553
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 17 deletions.
diff --git a/cloudprober.go b/cloudprober.go
@@ -68,11 +68,10 @@ var cloudProber struct {
 	prober          *prober.Prober
 	defaultServerLn net.Listener
 	defaultGRPCLn   net.Listener
-	rawConfig       string
-	parsedConfig    string
+	configSource    config.ConfigSource
 	config          *configpb.ProberConfig
 	cancelInitCtx   context.CancelFunc
-	sync.Mutex
+	sync.RWMutex
 }
 
 func getServerHost(c *configpb.ProberConfig) string {
@@ -227,8 +226,7 @@ func InitWithConfigSource(configSrc config.ConfigSource) error {
 
 	cloudProber.prober = pr
 	cloudProber.config = cfg
-	cloudProber.rawConfig = configSrc.RawConfig()
-	cloudProber.parsedConfig = configSrc.ParsedConfig()
+	cloudProber.configSource = configSrc
 	cloudProber.defaultServerLn = ln
 	cloudProber.defaultGRPCLn = grpcLn
 	cloudProber.cancelInitCtx = cancelFunc
@@ -258,9 +256,8 @@ func Start(ctx context.Context) {
 		defer cloudProber.Unlock()
 		cloudProber.defaultServerLn = nil
 		cloudProber.defaultGRPCLn = nil
-		cloudProber.rawConfig = ""
-		cloudProber.parsedConfig = ""
 		cloudProber.config = nil
+		cloudProber.configSource = nil
 		cloudProber.prober = nil
 	}()
 
@@ -281,28 +278,34 @@ func Start(ctx context.Context) {
 
 // GetConfig returns the prober config.
 func GetConfig() *configpb.ProberConfig {
-	cloudProber.Lock()
-	defer cloudProber.Unlock()
+	cloudProber.RLock()
+	defer cloudProber.RUnlock()
 	return cloudProber.config
 }
 
 // GetRawConfig returns the prober config in text proto format.
 func GetRawConfig() string {
-	cloudProber.Lock()
-	defer cloudProber.Unlock()
-	return cloudProber.rawConfig
+	cloudProber.RLock()
+	defer cloudProber.RUnlock()
+	return cloudProber.configSource.RawConfig()
 }
 
 // GetParsedConfig returns the parsed prober config.
 func GetParsedConfig() string {
-	cloudProber.Lock()
-	defer cloudProber.Unlock()
-	return cloudProber.parsedConfig
+	cloudProber.RLock()
+	defer cloudProber.RUnlock()
+	return cloudProber.configSource.ParsedConfig()
 }
 
 // GetInfo returns information on all the probes, servers and surfacers.
 func GetInfo() (map[string]*probes.ProbeInfo, []*surfacers.SurfacerInfo, []*servers.ServerInfo) {
-	cloudProber.Lock()
-	defer cloudProber.Unlock()
+	cloudProber.RLock()
+	defer cloudProber.RUnlock()
 	return cloudProber.prober.Probes, cloudProber.prober.Surfacers, cloudProber.prober.Servers
 }
+
+func GetProber() *prober.Prober {
+	cloudProber.RLock()
+	defer cloudProber.RUnlock()
+	return cloudProber.prober
+}
diff --git a/cloudprober_test.go b/cloudprober_test.go
@@ -21,6 +21,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/cloudprober/cloudprober/config"
 	configpb "github.com/cloudprober/cloudprober/config/proto"
 	serverspb "github.com/cloudprober/cloudprober/internal/servers/proto"
 	udpserverpb "github.com/cloudprober/cloudprober/internal/servers/udp/proto"
@@ -30,6 +31,7 @@ import (
 	"github.com/cloudprober/cloudprober/surfacers"
 	surfacerspb "github.com/cloudprober/cloudprober/surfacers/proto"
 	targetspb "github.com/cloudprober/cloudprober/targets/proto"
+	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 )
@@ -217,3 +219,41 @@ func TestRestart(t *testing.T) {
 		})
 	}
 }
+
+func TestCloudproberConfig(t *testing.T) {
+	rawCfg := `probe { type: PING, name: "test_probe", targets { host_names: "localhost" }}`
+	f, err := os.CreateTemp("", "cloudprober_test")
+	if err != nil {
+		t.Fatalf("os.CreateTemp(): %v", err)
+	}
+	defer os.Remove(f.Name())
+	os.WriteFile(f.Name(), []byte(rawCfg), 0644)
+
+	tests := []struct {
+		name             string
+		fileName         string
+		wantProbename    string
+		wantRawConfig    string
+		wantParsedConfig string
+	}{
+		{
+			name:             "config from file",
+			fileName:         f.Name(),
+			wantProbename:    "test_probe",
+			wantRawConfig:    rawCfg,
+			wantParsedConfig: rawCfg,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			configSrc := config.ConfigSourceWithFile(tt.fileName)
+			cloudProber.configSource = configSrc
+			cloudProber.config, _ = configSrc.GetConfig()
+
+			assert.Equal(t, tt.wantProbename, GetConfig().GetProbe()[0].GetName(), "GetConfig()")
+			assert.Equal(t, tt.wantRawConfig, GetRawConfig(), "GetRawConfig()")
+			assert.Equal(t, tt.wantParsedConfig, GetParsedConfig(), "GetParsedConfig()")
+		})
+	}
+}