Small restructure of plugin logic (#2221)

node/node.go

@@ -22,7 +22,7 @@ import (
 	"github.com/NethermindEth/juno/l1"
 	"github.com/NethermindEth/juno/migration"
 	"github.com/NethermindEth/juno/p2p"
-	junoplugin "github.com/NethermindEth/juno/plugin"
+	"github.com/NethermindEth/juno/plugin"
 	"github.com/NethermindEth/juno/rpc"
 	"github.com/NethermindEth/juno/service"
 	adaptfeeder "github.com/NethermindEth/juno/starknetdata/feeder"
@@ -159,15 +159,13 @@ func New(cfg *Config, version string) (*Node, error) { //nolint:gocyclo,funlen
 	synchronizer := sync.New(chain, adaptfeeder.New(client), log, cfg.PendingPollInterval, dbIsRemote)
 	gatewayClient := gateway.NewClient(cfg.Network.GatewayURL, log).WithUserAgent(ua).WithAPIKey(cfg.GatewayAPIKey)
 
-	pluginService := junoplugin.New(log)
 	if cfg.PluginPath != "" {
-		plugin, err := junoplugin.Load(cfg.PluginPath)
+		p, err := plugin.Load(cfg.PluginPath)
 		if err != nil {
 			return nil, err
 		}
-		synchronizer.WithPlugin(plugin)
-		pluginService.WithPlugin(plugin)
-		services = append(services, pluginService)
+		synchronizer.WithPlugin(p)
+		services = append(services, plugin.NewService(p))
 	}
 
 	var p2pService *p2p.Service

plugin/plugin.go

@@ -1,43 +1,13 @@
-package junoplugin
+package plugin
 
 import (
-	"context"
 	"fmt"
-	"plugin"
-	"sync"
+	stdplugin "plugin"
 
 	"github.com/NethermindEth/juno/core"
 	"github.com/NethermindEth/juno/core/felt"
-	"github.com/NethermindEth/juno/utils"
 )
 
-type PluginService struct {
-	jPlugin JunoPlugin
-	wg      sync.WaitGroup
-	log     utils.SimpleLogger
-}
-
-func New(log utils.SimpleLogger) *PluginService {
-	return &PluginService{wg: sync.WaitGroup{}, log: log}
-}
-
-func (p *PluginService) WithPlugin(jPlugin JunoPlugin) {
-	p.jPlugin = jPlugin
-}
-
-func (p *PluginService) Run(ctx context.Context) error {
-	p.wg.Add(1)
-	go func() {
-		defer p.wg.Done()
-		<-ctx.Done()
-		if err := p.jPlugin.Shutdown(); err != nil {
-			p.log.Errorw("Error while calling plugin Shutdown() function", "err", err)
-		}
-	}()
-	p.wg.Wait()
-	return nil
-}
-
 //go:generate mockgen -destination=../mocks/mock_plugin.go -package=mocks github.com/NethermindEth/juno/plugin JunoPlugin
 type JunoPlugin interface {
 	Init() error
@@ -54,7 +24,7 @@ type BlockAndStateUpdate struct {
 }
 
 func Load(pluginPath string) (JunoPlugin, error) {
-	plug, err := plugin.Open(pluginPath)
+	plug, err := stdplugin.Open(pluginPath)
 	if err != nil {
 		return nil, fmt.Errorf("error loading plugin .so file: %w", err)
 	}

plugin/plugin_test.go

@@ -1,4 +1,4 @@
-package junoplugin_test
+package plugin_test
 
 import (
 	"context"

plugin/service.go

@@ -0,0 +1,19 @@
+package plugin
+
+import "context"
+
+// Service provides an abstraction for signalling the plugin to shut down.
+type Service struct {
+	plugin JunoPlugin
+}
+
+func NewService(plugin JunoPlugin) *Service {
+	return &Service{
+		plugin: plugin,
+	}
+}
+
+func (p *Service) Run(ctx context.Context) error {
+	<-ctx.Done()
+	return p.plugin.Shutdown()
+}

plugin/service_test.go

@@ -0,0 +1,45 @@
+package plugin_test
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/NethermindEth/juno/mocks"
+	"github.com/NethermindEth/juno/plugin"
+	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
+)
+
+func TestService(t *testing.T) {
+	t.Run("shutdown ok", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
+		p := mocks.NewMockJunoPlugin(ctrl)
+		p.EXPECT().Shutdown().Return(nil)
+		service := plugin.NewService(p)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		// after ^ this ctx already cancelled
+
+		err := service.Run(ctx)
+		require.NoError(t, err)
+	})
+	t.Run("shutdown with error", func(t *testing.T) {
+		ctrl := gomock.NewController(t)
+
+		shutdownErr := errors.New("error during shutdown")
+
+		p := mocks.NewMockJunoPlugin(ctrl)
+		p.EXPECT().Shutdown().Return(shutdownErr)
+		service := plugin.NewService(p)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+		// after ^ this ctx already cancelled
+
+		err := service.Run(ctx)
+		require.Equal(t, shutdownErr, err)
+	})
+}