pkg/plug/grpc.go

@@ -61,6 +61,7 @@ func (s *server) TaskStart(ctx context.Context, creation *gRPC.PlugTaskCreation)
 	// by using this method we assign a new context to each new request we get.
 	// This can be used for example to close the context with the given id
 	ctx = trace.ContextWithSpanContext(context.Background(), trace.NewSpanContext(trace.SpanContextConfig{TraceID: span.SpanContext().TraceID()}))
+
 	taskCtx, cancel := context.WithCancel(ctx)
 	s.ctx[id] = cancel
 	span.AddEvent("Created new context for task", trace.WithAttributes(attribute.String("task_id", id)))

pkg/plug/otlp.go

@@ -1,31 +0,0 @@
-package plug
-
-import (
-	"context"
-	"git.anthrove.art/Anthrove/plug-sdk/v3/pkg/telemetry"
-	log "github.com/sirupsen/logrus"
-	"go.opentelemetry.io/otel/codes"
-)
-
-func SetupOpenTelemtry(ctx context.Context, serviceName string) error {
-	ctx, span := tracer.Start(ctx, "SetupOpenTelemtry")
-	defer span.End()
-
-	err := telemetry.SetupMeterProvider(serviceName)
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, err.Error())
-		log.WithContext(ctx).WithError(err).Error("Failed to setup meter provider")
-		return err
-	}
-
-	err = telemetry.SetupTraceProvider(ctx, serviceName)
-	if err != nil {
-		span.RecordError(err)
-		span.SetStatus(codes.Error, err.Error())
-		log.WithContext(ctx).WithError(err).Error("Failed to setup trace provider")
-		return err
-	}
-
-	return nil
-}

pkg/plug/server.go

@@ -2,11 +2,13 @@ package plug
 
 import (
 	"context"
+	"fmt"
 	"net"
 
 	"git.anthrove.art/Anthrove/otter-space-sdk/v3/pkg/database"
 	"git.anthrove.art/Anthrove/otter-space-sdk/v3/pkg/models"
 	pb "git.anthrove.art/Anthrove/plug-sdk/v3/pkg/grpc"
+	"git.anthrove.art/Anthrove/plug-sdk/v3/pkg/telemetry"
 	log "github.com/sirupsen/logrus"
 	"go.opentelemetry.io/otel"
 	"go.opentelemetry.io/otel/attribute"
@@ -27,40 +29,57 @@ type TaskExecution func(ctx context.Context, userSource models.UserSource, deepS
 type SendMessageExecution func(ctx context.Context, userSource models.UserSource, message string) error
 type GetMessageExecution func(ctx context.Context, userSource models.UserSource) ([]Message, error)
 
-var (
+type Plug struct {
+	address               string
+	port                  string
+	ctx                   context.Context
 	taskExecutionFunction TaskExecution
 	sendMessageExecution  SendMessageExecution
 	getMessageExecution   GetMessageExecution
-)
+	source                models.Source
+	plugName              string
+}
 
-func Listen(ctx context.Context, listenAddr string, source models.Source) error {
-	ctx, span := tracer.Start(ctx, "Listen")
+func NewPlug(ctx context.Context, plugName string, address string, port string, source models.Source) Plug {
+	return Plug{
+		ctx:      ctx,
+		address:  address,
+		port:     port,
+		source:   source,
+		plugName: plugName,
+	}
+}
+
+func (p *Plug) Listen() error {
+	ctx, span := tracer.Start(p.ctx, "Listen")
 	defer span.End()
 
 	var err error
+	var source models.Source
 
 	span.SetAttributes(
-		attribute.String("source_display_name", string(source.DisplayName)),
-		attribute.String("source_domain", string(source.Domain)),
+		attribute.String("source_display_name", string(p.source.DisplayName)),
+		attribute.String("source_domain", string(p.source.Domain)),
 	)
 
 	sourceFields := log.Fields{
-		"source_display_name": source.DisplayName,
-		"source_domain":       source.Domain,
+		"source_display_name": p.source.DisplayName,
+		"source_domain":       p.source.Domain,
 	}
 
 	serverFields := log.Fields{
-		"address": listenAddr,
+		"address": p.address,
+		"port":    p.port,
 	}
 
-	source, err = database.GetSourceByDomain(ctx, source.Domain)
+	source, err = database.GetSourceByDomain(ctx, p.source.Domain)
 	if err != nil {
 		if err.Error() == "Database error: NoDataFound" {
 			span.AddEvent("No Source found, initializing source")
 
 			log.WithContext(ctx).WithFields(sourceFields).Info("No Source found, initializing source!")
 
-			source, err = database.CreateSource(ctx, source)
+			source, err = database.CreateSource(ctx, p.source)
 			if err != nil {
 				span.RecordError(err)
 				span.SetStatus(codes.Error, err.Error())
@@ -79,11 +98,13 @@ func Listen(ctx context.Context, listenAddr string, source models.Source) error
 		}
 	}
 
+	p.source = source
+
 	span.SetAttributes(
-		attribute.String("source_id", string(source.ID)),
+		attribute.String("source_id", string(p.source.ID)),
 	)
 
-	lis, err := net.Listen("tcp", listenAddr)
+	lis, err := net.Listen("tcp", fmt.Sprintf("%s:%s", p.address, p.port))
 	if err != nil {
 		span.RecordError(err)
 		span.SetStatus(codes.Error, err.Error())
@@ -93,34 +114,70 @@ func Listen(ctx context.Context, listenAddr string, source models.Source) error
 
 	grpcServer := grpc.NewServer()
 
-	pb.RegisterPlugConnectorServer(grpcServer, NewGrpcServer(source, taskExecutionFunction, sendMessageExecution, getMessageExecution))
+	pb.RegisterPlugConnectorServer(grpcServer, NewGrpcServer(p.source, p.taskExecutionFunction, p.sendMessageExecution, p.getMessageExecution))
 
-	go func() {
 	err = grpcServer.Serve(lis)
 	if err != nil {
 		span.RecordError(err)
 		span.SetStatus(codes.Error, err.Error())
 		log.WithContext(ctx).WithError(err).Error("Failed to serve gRPC server")
-		}
-	}()
-
-	select {
-	case <-ctx.Done():
-		log.WithContext(ctx).Info("Context stopped! Shutdown Plug")
-		grpcServer.GracefulStop()
+		return err
 	}
 
 	return nil
 }
 
-func SetTaskExecutionFunction(function TaskExecution) {
-	taskExecutionFunction = function
+func (p *Plug) GetSource() models.Source {
+	_, span := tracer.Start(p.ctx, "GetSource")
+	defer span.End()
+
+	span.SetAttributes(attribute.String("source_domain", string(p.source.Domain)))
+	return p.source
 }
 
-func SetSendMessageExecutionFunction(function SendMessageExecution) {
-	sendMessageExecution = function
+func (p *Plug) TaskExecutionFunction(function TaskExecution) {
+	_, span := tracer.Start(p.ctx, "TaskExecutionFunction")
+	defer span.End()
+
+	p.taskExecutionFunction = function
+	span.AddEvent("Task execution function set")
 }
 
-func SetGetMessageExecutionFunction(function GetMessageExecution) {
-	getMessageExecution = function
+func (p *Plug) SendMessageExecution(function SendMessageExecution) {
+	_, span := tracer.Start(p.ctx, "SendMessageExecution")
+	defer span.End()
+
+	p.sendMessageExecution = function
+	span.AddEvent("Send message execution function set")
+}
+
+func (p *Plug) GetMessageExecution(function GetMessageExecution) {
+	_, span := tracer.Start(p.ctx, "GetMessageExecution")
+	defer span.End()
+
+	p.getMessageExecution = function
+	span.AddEvent("Get message execution function set")
+}
+
+func (p Plug) SetupOpenTelemtry() error {
+	ctx, span := tracer.Start(p.ctx, "SetupOpenTelemtry")
+	defer span.End()
+
+	err := telemetry.SetupMeterProvider(p.plugName)
+	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
+		log.WithContext(ctx).WithError(err).Error("Failed to setup meter provider")
+		return err
+	}
+
+	err = telemetry.SetupTraceProvider(ctx, p.plugName)
+	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, err.Error())
+		log.WithContext(ctx).WithError(err).Error("Failed to setup trace provider")
+		return err
+	}
+
+	return nil
 }