[SDK] Refine StageLogPersister handling in SDK

pkg/app/pipedv1/plugin/wait/wait.go

@@ -22,7 +22,6 @@
 
 	"go.uber.org/zap"
 
-	"github.com/pipe-cd/pipecd/pkg/app/piped/logpersister"
 	"github.com/pipe-cd/pipecd/pkg/plugin/sdk"
 )
 
@@ -35,7 +34,7 @@
 func (p *plugin) executeWait(ctx context.Context, in *sdk.ExecuteStageInput) sdk.StageStatus {
 	opts, err := decode(in.Request.StageConfig)
 	if err != nil {
-		in.Client.LogPersister.Errorf("failed to decode the stage config: %v", err)
+		in.Client.LogPersister().Errorf("failed to decode the stage config: %v", err)
 		return sdk.StageStatusFailure
 	}
 
@@ -49,10 +48,10 @@
 	}
 	p.saveStartTime(ctx, in.Client, initialStart, in.Logger)
 
-	return wait(ctx, duration, initialStart, in.Client.LogPersister)
+	return wait(ctx, duration, initialStart, in.Client.LogPersister())
 }
 
-func wait(ctx context.Context, duration time.Duration, initialStart time.Time, slp logpersister.StageLogPersister) sdk.StageStatus {
+func wait(ctx context.Context, duration time.Duration, initialStart time.Time, slp sdk.StageLogPersister) sdk.StageStatus {
 	remaining := duration - time.Since(initialStart)
 	if remaining <= 0 {
 		// When this stage restarted and the duration has already passed.

pkg/plugin/sdk/client.go

@@ -17,7 +17,6 @@
 import (
 	"context"
 
-	"github.com/pipe-cd/pipecd/pkg/plugin/logpersister"
 	"github.com/pipe-cd/pipecd/pkg/plugin/pipedapi"
 	"github.com/pipe-cd/pipecd/pkg/plugin/pipedservice"
 )
@@ -40,8 +39,21 @@
 	// This field exists only when the client is working with a specific stage; for example, when this client is passed as the ExecuteStage method's argument.
 	stageID string
 
-	// TODO: Define another interface (e.g. Remove Complete() method)
-	LogPersister logpersister.StageLogPersister
+	// logPersister is used to persist the stage logs.
+	// This field exists only when the client is working with a specific stage; for example, when this client is passed as the ExecuteStage method's argument.
+	logPersister StageLogPersister
+}
+
+// StageLogPersister is a interface for persisting the stage logs.
+// Use this to persist the stage logs and make it viewable on the UI.
+type StageLogPersister interface {
+	Write(log []byte) (int, error)
+	Info(log string)
+	Infof(format string, a ...interface{})
+	Success(log string)
+	Successf(format string, a ...interface{})
+	Error(log string)
+	Errorf(format string, a ...interface{})
 }
 
 // GetStageMetadata gets the metadata of the current stage.
@@ -118,3 +130,12 @@
 	})
 	return resp.Value, err
 }
+
+// LogPersister returns the stage log persister.
+// Use this to persist the stage logs and make it viewable on the UI.
+// This method should be called only when the client is working with a specific stage, for example, when this client is passed as the ExecuteStage method's argument.
+// Otherwise, it will return nil.
+// TODO: we should consider returning an error instead of nil, or return logger which prints to stdout.
+func (c *Client) LogPersister() StageLogPersister {
+	return c.logPersister
+}

pkg/plugin/sdk/deployment.go

@@ -30,6 +30,7 @@
 	"github.com/pipe-cd/pipecd/pkg/plugin/api/v1alpha1/deployment"
 	"github.com/pipe-cd/pipecd/pkg/plugin/logpersister"
 	"github.com/pipe-cd/pipecd/pkg/plugin/pipedapi"
+	"github.com/pipe-cd/pipecd/pkg/plugin/signalhandler"
 )
 
 var (
@@ -170,14 +171,24 @@
 func (s *DeploymentPluginServiceServer[Config, DeployTargetConfig]) BuildQuickSyncStages(context.Context, *deployment.BuildQuickSyncStagesRequest) (*deployment.BuildQuickSyncStagesResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method BuildQuickSyncStages not implemented")
 }
-func (s *DeploymentPluginServiceServer[Config, DeployTargetConfig]) ExecuteStage(ctx context.Context, request *deployment.ExecuteStageRequest) (*deployment.ExecuteStageResponse, error) {
+func (s *DeploymentPluginServiceServer[Config, DeployTargetConfig]) ExecuteStage(ctx context.Context, request *deployment.ExecuteStageRequest) (response *deployment.ExecuteStageResponse, _ error) {
+	lp := s.logPersister.StageLogPersister(request.GetInput().GetDeployment().GetId(), request.GetInput().GetStage().GetId())
+	defer func() {
+		// When termination signal received and the stage is not completed yet, we should not mark the log persister as completed.
+		// This can occur when the piped is shutting down while the stage is still running.
+		if !response.GetStatus().IsCompleted() && signalhandler.Terminated() {
+			return
+		}
+		lp.Complete(time.Minute)
+	}()
+
 	client := &Client{
 		base:          s.client,
 		pluginName:    s.Name(),
 		applicationID: request.GetInput().GetDeployment().GetApplicationId(),
 		deploymentID:  request.GetInput().GetDeployment().GetId(),
 		stageID:       request.GetInput().GetStage().GetId(),
-		LogPersister:  s.logPersister.StageLogPersister(request.GetInput().GetDeployment().GetId(), request.GetInput().GetStage().GetId()),
+		logPersister:  lp,
 	}
 	return executeStage(ctx, s.base, &s.config, nil, client, request, s.logger) // TODO: pass the deployTargets
 }
@@ -240,15 +251,26 @@
 func (s *PipelineSyncPluginServiceServer[Config, DeployTargetConfig]) BuildQuickSyncStages(context.Context, *deployment.BuildQuickSyncStagesRequest) (*deployment.BuildQuickSyncStagesResponse, error) {
 	return nil, status.Errorf(codes.Unimplemented, "method BuildQuickSyncStages not implemented")
 }
-func (s *PipelineSyncPluginServiceServer[Config, DeployTargetConfig]) ExecuteStage(ctx context.Context, request *deployment.ExecuteStageRequest) (*deployment.ExecuteStageResponse, error) {
+func (s *PipelineSyncPluginServiceServer[Config, DeployTargetConfig]) ExecuteStage(ctx context.Context, request *deployment.ExecuteStageRequest) (response *deployment.ExecuteStageResponse, _ error) {
+	lp := s.logPersister.StageLogPersister(request.GetInput().GetDeployment().GetId(), request.GetInput().GetStage().GetId())
+	defer func() {
+		// When termination signal received and the stage is not completed yet, we should not mark the log persister as completed.
+		// This can occur when the piped is shutting down while the stage is still running.
+		if !response.GetStatus().IsCompleted() && signalhandler.Terminated() {
+			return
+		}
+		lp.Complete(time.Minute)
+	}()
+
 	client := &Client{
 		base:          s.client,
 		pluginName:    s.Name(),
 		applicationID: request.GetInput().GetDeployment().GetApplicationId(),
 		deploymentID:  request.GetInput().GetDeployment().GetId(),
 		stageID:       request.GetInput().GetStage().GetId(),
-		LogPersister:  s.logPersister.StageLogPersister(request.GetInput().GetDeployment().GetId(), request.GetInput().GetStage().GetId()),
+		logPersister:  lp,
 	}
+
 	return executeStage(ctx, s.base, &s.config, nil, client, request, s.logger) // TODO: pass the deployTargets
 }