Implement ConfigReload on full config SET update

common_utils/notification_producer.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"strings"
+	"time"
 
 	sdcfg "github.com/sonic-net/sonic-gnmi/sonic_db_config"
 
@@ -44,6 +45,23 @@ func GetRedisDBClient() (*redis.Client, error) {
 	return rclient, nil
 }
 
+func IsAutoRestartEnabled() bool {
+	redisDb, err := GetRedisDBClient()
+	if err != nil {
+		log.Errorf("failed to connect to Redis for auto_restart check: %v", err)
+		return false
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
+	defer cancel()
+	value, err := redisDb.HGet(ctx, "FEATURE|gnmi", "auto_restart").Result()
+	if err != nil {
+		log.Errorf("failed to read FEATURE|gnmi auto_restart: %v", err)
+		return false
+	}
+	return value == "enabled"
+}
+
 // NotificationProducer provides utilities for sending messages using notification channel.
 // NewNotificationProducer must be called for a new producer.
 // Close must be called when finished.

gnmi_server/client_subscribe.go

@@ -5,6 +5,7 @@ import (
 	"github.com/Workiva/go-datastructures/queue"
 	log "github.com/golang/glog"
 	gnmipb "github.com/openconfig/gnmi/proto/gnmi"
+	spb "github.com/sonic-net/sonic-gnmi/proto"
 	sdc "github.com/sonic-net/sonic-gnmi/sonic_data_client"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -44,6 +45,9 @@ type Client struct {
 	logLevel int
 }
 
+// key used to forcefully shutdown all active gNMI Subscribes
+const ConfigReloadRequested string = "full config reload requested"
+
 // Syslog level for error
 const logLevelError int = 3
 const logLevelDebug int = 7
@@ -252,6 +256,9 @@ func (c *Client) Run(stream gnmipb.GNMI_SubscribeServer, config *Config) (err er
 	go c.recv(stream)
 	err = c.send(stream, dc)
 	c.Close()
+	if err != nil && err.Error() == ConfigReloadRequested {
+		return grpc.Errorf(codes.Unavailable, "%s", err)
+	}
 	// Wait until all child go routines exited
 	c.w.Wait()
 	return grpc.Errorf(codes.InvalidArgument, "%s", err)
@@ -281,6 +288,27 @@ func (c *Client) Close() {
 	}
 }
 
+// Abort() enqueues a Fatal message to the subscribed client's queue
+// (if not already disposed). This will cause c.send() to exit out
+// and proceed with client connection teardown.
+//
+// If calling to abort subscribe RPCs in preparation for a config reload,
+// pass in `ConfigReloadRequested` as the argument.
+func (c *Client) Abort(reason string) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if c.q == nil || c.q.Disposed() {
+		return
+	}
+
+	c.q.Put(sdc.Value{
+		Value: &spb.Value{
+			Fatal: reason,
+		},
+	})
+}
+
 func (c *Client) recv(stream gnmipb.GNMI_SubscribeServer) {
 	defer c.Close()
 

gnmi_server/server.go

@@ -825,6 +825,14 @@ func authenticate(config *Config, ctx context.Context, target string, writeAcces
 	return ctx, nil
 }
 
+func (s *Server) dropSubscribeConns(msg string) {
+	s.cMu.Lock()
+	defer s.cMu.Unlock()
+	for _, c := range s.clients {
+		c.Abort(msg)
+	}
+}
+
 // Subscribe implements the gNMI Subscribe RPC.
 func (s *Server) Subscribe(stream gnmipb.GNMI_SubscribeServer) error {
 	ctx := stream.Context()
@@ -1052,6 +1060,18 @@ func SaveOnSetEnabled() error {
 // SaveOnSetDisabeld does nothing.
 func saveOnSetDisabled() error { return nil }
 
+func (s *Server) PostResCallback(dc *sdc.MixedDbClient) {
+	autoRestartEnabled := common_utils.IsAutoRestartEnabled()
+	if autoRestartEnabled {
+		log.V(1).Info("server: waiting for in-flight RPCs to finish")
+		s.dropSubscribeConns(ConfigReloadRequested)
+		s.Stop() // blocking
+	} else {
+		log.Warning("gNMI auto-restart disabled, skipping RPC aborts.")
+	}
+	dc.MaybeRunCallback(autoRestartEnabled)
+}
+
 func (s *Server) Set(ctx context.Context, req *gnmipb.SetRequest) (*gnmipb.SetResponse, error) {
 	e := s.ReqFromMaster(req, &s.masterEID)
 	if e != nil {
@@ -1192,6 +1212,10 @@ func (s *Server) Set(ctx context.Context, req *gnmipb.SetRequest) (*gnmipb.SetRe
 		common_utils.IncCounter(common_utils.GNMI_SET_FAIL)
 	} else {
 		s.SaveStartupConfig()
+
+		if mdc, ok := dc.(*sdc.MixedDbClient); ok && mdc.ConfigReloadRequested() {
+			go s.PostResCallback(mdc)
+		}
 	}
 
 	return &gnmipb.SetResponse{

gnmi_server/server_test.go

@@ -36,6 +36,7 @@ import (
 	"github.com/sonic-net/sonic-gnmi/swsscommon"
 	"github.com/sonic-net/sonic-gnmi/test_utils"
 	testcert "github.com/sonic-net/sonic-gnmi/testdata/tls"
+	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc/security/advancedtls"
 
 	"github.com/golang/protobuf/proto"
@@ -5546,6 +5547,31 @@ print('%s')
 	s.Stop()
 }
 
+func TestGNMIPostResCallback(t *testing.T) {
+	sdcfg.Init()
+	s := createServer(t, 8080)
+	defer s.Stop()
+	callbackCount := 0
+	mock1 := gomonkey.ApplyFunc(common_utils.IsAutoRestartEnabled, func() bool { return true })
+	defer mock1.Reset()
+	mock2 := gomonkey.ApplyMethod(reflect.TypeOf(&sdc.MixedDbClient{}), "MaybeRunCallback", func(dc *sdc.MixedDbClient, restartEnabled bool) {
+		callbackCount++
+	})
+	defer mock2.Reset()
+	mock3 := gomonkey.ApplyFunc((*Server).dropSubscribeConns, func(s *Server, reason string) {})
+	defer mock3.Reset()
+	mock4 := gomonkey.ApplyMethod(reflect.TypeOf(&Server{}), "Stop", func(s *Server) {})
+	defer mock4.Reset()
+
+	dc := &sdc.MixedDbClient{}
+	s.PostResCallback(dc)
+	assert.True(t, callbackCount == 1)
+
+	mock1 = gomonkey.ApplyFunc(common_utils.IsAutoRestartEnabled, func() bool { return false })
+	s.PostResCallback(dc)
+	assert.True(t, callbackCount == 2)
+}
+
 func TestGNMINative(t *testing.T) {
 	mock1 := gomonkey.ApplyFunc(dbus.SystemBus, func() (conn *dbus.Conn, err error) {
 		return &dbus.Conn{}, nil
@@ -5562,6 +5588,16 @@ func TestGNMINative(t *testing.T) {
 	defer mock2.Reset()
 	mock3 := gomonkey.ApplyFunc(sdc.RunPyCode, func(text string) error { return nil })
 	defer mock3.Reset()
+	mock4 := gomonkey.ApplyMethod(reflect.TypeOf(&ssc.DbusClient{}), "ConfigReload", func(c *ssc.DbusClient, config string) error {
+		common_utils.IncCounter(common_utils.DBUS_CONFIG_RELOAD)
+		return nil
+	})
+	defer mock4.Reset()
+	mock5 := gomonkey.ApplyMethod(reflect.TypeOf(&Server{}), "PostResCallback", func(s *Server, dc *sdc.MixedDbClient) {
+		assert.True(t, dc.ConfigReloadRequested(), "ConfigReload callback should have been registered")
+		s.dropSubscribeConns(ConfigReloadRequested)
+	})
+	defer mock5.Reset()
 
 	sdcfg.Init()
 	s := createServer(t, 8080)

sonic_data_client/config_reload_test.go

@@ -0,0 +1,175 @@
+package client
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+
+	"github.com/agiledragon/gomonkey/v2"
+	gnmipb "github.com/openconfig/gnmi/proto/gnmi"
+	"github.com/sonic-net/sonic-gnmi/common_utils"
+	ssc "github.com/sonic-net/sonic-gnmi/sonic_service_client"
+	"github.com/stretchr/testify/assert"
+)
+
+func getFullConfigUpdateMessage(jsonContent string) []*gnmipb.Update {
+	return []*gnmipb.Update{{
+		Val: &gnmipb.TypedValue{
+			Value: &gnmipb.TypedValue_JsonIetfVal{
+				JsonIetfVal: []byte(jsonContent),
+			},
+		},
+	}}
+}
+
+func TestSetFullConfig(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	mock := gomonkey.ApplyFunc(RunPyCode, func(text string) error {
+		return nil
+	})
+	defer mock.Reset()
+
+	mock2 := gomonkey.ApplyMethod(reflect.TypeOf(&ssc.DbusClient{}), "ConfigReload", func(dbus *ssc.DbusClient, config string) error {
+		common_utils.IncCounter(common_utils.DBUS_CONFIG_RELOAD)
+		return nil
+	})
+	defer mock2.Reset()
+
+	c := &MixedDbClient{workPath: tmpDir}
+	update := getFullConfigUpdateMessage(`{"DEVICE_METADATA": {}}`)
+
+	// technically we should pass in a delete path as well,
+	// but that's only for determining the specific config
+	// operation to invoke from the SetConfigDB() level.
+	// delete isn't used internally in SetFullConfig
+	err := c.SetFullConfig(nil, nil, update)
+	if err != nil {
+		t.Fatalf("expected no error, got: %v", err)
+	}
+	if !c.ConfigReloadRequested() {
+		t.Fatal("expected configReloadCallback to be registered after successful SetFullConfig")
+	}
+
+	var counts [int(common_utils.COUNTER_SIZE)]uint64
+	preCount := counts[int(common_utils.DBUS_CONFIG_RELOAD)]
+	if err = common_utils.GetMemCounters(&counts); err != nil {
+		t.Fatalf("failed to get memory counters: %v", err)
+	}
+	c.MaybeRunCallback(true)
+	if err = common_utils.GetMemCounters(&counts); err != nil {
+		t.Fatalf("failed to get memory counters: %v", err)
+	}
+
+	actualCount := counts[int(common_utils.DBUS_CONFIG_RELOAD)]
+	if actualCount != preCount+1 {
+		t.Fatalf("expected DBUS_CONFIG_RELOAD to be %d, got: %v", preCount+1, actualCount)
+	}
+}
+
+func TestSetFullConfigDbusClientFail(t *testing.T) {
+	mock1 := gomonkey.ApplyFunc(ssc.NewDbusClient, func() (ssc.Service, error) {
+		return nil, fmt.Errorf("failed to create dbus client")
+	})
+	defer mock1.Reset()
+
+	configReloadCalled := false
+	mock2 := gomonkey.ApplyMethod(reflect.TypeOf(&ssc.DbusClient{}), "ConfigReload", func(dbus *ssc.DbusClient, config string) error {
+		configReloadCalled = true
+		return nil
+	})
+	defer mock2.Reset()
+
+	mock3 := gomonkey.ApplyFunc(RunPyCode, func(text string) error {
+		return nil
+	})
+	defer mock3.Reset()
+
+	tmpDir := t.TempDir()
+	c := &MixedDbClient{workPath: tmpDir}
+	update := getFullConfigUpdateMessage(`{"DEVICE_METADATA": {}}`)
+
+	err := c.SetFullConfig(nil, nil, update)
+	if err != nil {
+		t.Fatalf("expected no error, got: %v", err)
+	}
+	if !c.ConfigReloadRequested() {
+		t.Fatal("expected configReloadCallback to be registered after successful SetFullConfig")
+	}
+
+	c.MaybeRunCallback(true)
+	assert.False(t, configReloadCalled)
+}
+
+func TestSetFullConfigReloadFail(t *testing.T) {
+	mock := gomonkey.ApplyFunc(RunPyCode, func(text string) error {
+		return nil
+	})
+	defer mock.Reset()
+
+	mock2 := gomonkey.ApplyMethod(reflect.TypeOf(&ssc.DbusClient{}), "ConfigReload", func(dbus *ssc.DbusClient, config string) error {
+		return fmt.Errorf("config reload failed")
+	})
+	defer mock2.Reset()
+
+	verifyClientSetup := func() *MixedDbClient {
+		tmpDir := t.TempDir()
+
+		c := &MixedDbClient{workPath: tmpDir}
+		update := getFullConfigUpdateMessage(`{"DEVICE_METADATA": {}}`)
+
+		err := c.SetFullConfig(nil, nil, update)
+		if err != nil {
+			t.Fatalf("expected no error, got: %v", err)
+		}
+		if !c.ConfigReloadRequested() {
+			t.Fatal("expected configReloadCallback to be registered after successful SetFullConfig")
+		}
+
+		return c
+	}
+
+	c := verifyClientSetup()
+	assert.NotPanics(t, func() { c.MaybeRunCallback(false) }) // gNMI auto restart disabled
+	c = verifyClientSetup()
+	assert.Panics(t, func() { c.MaybeRunCallback(true) }) // gNMI auto restart enabled
+}
+
+func TestSetFullConfigInvalidJson(t *testing.T) {
+	c := &MixedDbClient{}
+	update := getFullConfigUpdateMessage("")
+
+	err := c.SetFullConfig(nil, nil, update)
+	if err == nil {
+		t.Fatal("expected error for empty IETF JSON value")
+	}
+	if err.Error() != "Value encoding is not IETF JSON" {
+		t.Fatalf("unexpected error: '%v'", err)
+	}
+	if c.ConfigReloadRequested() {
+		t.Fatal("callback shouldn't be registered when JSON validation fails")
+	}
+}
+
+func TestSetFullConfigYangValidationFail(t *testing.T) {
+	tmpDir := t.TempDir()
+
+	mock := gomonkey.ApplyFunc(RunPyCode, func(text string) error {
+		return fmt.Errorf("YANG model mismatch")
+	})
+	defer mock.Reset()
+
+	c := &MixedDbClient{workPath: tmpDir}
+	update := getFullConfigUpdateMessage(`{"BAD_TABLE": {}}`)
+
+	err := c.SetFullConfig(nil, nil, update)
+	if err == nil {
+		t.Fatal("expected error when YANG validation fails")
+	}
+	if err.Error() != "Yang validation failed!" {
+		t.Fatalf("unexpected error message: %v", err)
+	}
+	if c.ConfigReloadRequested() {
+		t.Fatal("callback should NOT be registered when YANG validation fails")
+	}
+}