feat: #5 support self-upgrade command

cmd/ptah-agent/ptah-agent.go

@@ -6,6 +6,7 @@ import (
 	ptahAgent "github.com/ptah-sh/ptah-agent/internal/app/ptah-agent"
 	"log"
 	"os"
+	"path"
 	"strings"
 )
 
@@ -27,9 +28,24 @@ func main() {
 		log.Fatalln("PTAH_TOKEN is not set")
 	}
 
-	agent := ptahAgent.New(version, baseUrl, ptahToken)
+	rootDir := os.Getenv("PTAH_ROOT_DIR")
+	if rootDir == "" {
+		log.Fatalln("PTAH_ROOT_DIR is not set")
+	}
+
+	_, err := os.Stat(rootDir)
+	if err != nil {
+		log.Fatalln(err)
+	}
+
+	_, err = os.Stat(path.Join(rootDir, "versions"))
+	if err != nil {
+		log.Fatalln("versions dir not found:", err)
+	}
+
+	agent := ptahAgent.New(version, baseUrl, ptahToken, rootDir)
 
-	err := agent.Start(context.Background())
+	err = agent.Start(context.Background())
 	if err != nil {
 		log.Fatalln(err)
 	}

install.sh

@@ -150,9 +150,9 @@ set -e
 
 echo "Running in user space as \$(whoami)"
 
-SEED_VERSION="\$HOME/ptah-agent/versions/0.0.0"
+SEED_VERSION="\$HOME/ptah-agent/versions/v0.0.0"
 
-mkdir -p \$HOME/ptah-agent/versions/0.0.0
+mkdir -p \$HOME/ptah-agent/versions/v0.0.0
 
 curl -L https://github.com/ptah-sh/ptah-agent/releases/latest/download/ptah-agent-linux-x86_64.bin -o \$SEED_VERSION/ptah-agent
 
@@ -187,7 +187,7 @@ After=network.target
 [Service]
 User=$USER
 Group=$GROUP
-Environment=PTAH_HOME=/home/$USER/ptah-agent
+Environment=PTAH_ROOT_DIR=/home/$USER/ptah-agent
 Environment=PTAH_TOKEN=$PTAH_TOKEN
 Environment=PTAH_BASE_URL=$PTAH_BASE_URL
 Type=exec

internal/app/ptah-agent/agent_upgrade.go

@@ -0,0 +1,82 @@
+package ptah_agent
+
+import (
+	"context"
+	"fmt"
+	"github.com/pkg/errors"
+	t "github.com/ptah-sh/ptah-agent/internal/pkg/ptah-client"
+	"io"
+	"math"
+	"net/http"
+	"os"
+	"path"
+	"time"
+)
+
+func (e *taskExecutor) downloadAgentUpgrade(ctx context.Context, req *t.DownloadAgentUpgradeReq) (*t.DownloadAgentUpgradeRes, error) {
+	startTime := time.Now()
+
+	request, err := http.NewRequestWithContext(ctx, "GET", req.DownloadUrl, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	request.Header.Set("Accept", "application/octet-stream")
+
+	resp, err := http.DefaultClient.Do(request)
+	if err != nil {
+		return nil, err
+	}
+
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("download agent upgrade: unexpected status code %d", resp.StatusCode)
+	}
+
+	if resp.Header.Get("Content-Type") != "application/octet-stream" {
+		return nil, fmt.Errorf("download agent upgrade: unexpected content type %s", resp.Header.Get("Content-Type"))
+	}
+
+	file, err := os.Create(path.Join(e.rootDir, "versions", req.TargetVersion))
+	if err != nil {
+		return nil, err
+	}
+
+	defer file.Close()
+
+	written, err := io.Copy(file, resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	elapsedTime := time.Since(startTime).Seconds()
+
+	return &t.DownloadAgentUpgradeRes{DownloadTime: int(math.Ceil(elapsedTime)), FileSize: int(written)}, nil
+}
+
+func (e *taskExecutor) updateAgentSymlink(ctx context.Context, req *t.UpdateAgentSymlinkReq) (*t.UpdateAgentSymlinkRes, error) {
+	current := path.Join(e.rootDir, "current")
+
+	err := os.Remove(current)
+	if err != nil && !errors.Is(err, os.ErrNotExist) {
+		return nil, err
+	}
+
+	err = os.Symlink(path.Join(e.rootDir, "versions", req.TargetVersion), current)
+	if err != nil {
+		return nil, err
+	}
+
+	e.stop()
+
+	return &t.UpdateAgentSymlinkRes{}, nil
+}
+
+func (e *taskExecutor) confirmAgentUpgrade(ctx context.Context, req *t.ConfirmAgentUpgradeReq) (*t.ConfirmAgentUpgradeRes, error) {
+	if e.agent.Version == req.TargetVersion {
+		return &t.ConfirmAgentUpgradeRes{}, nil
+	}
+
+	return nil, fmt.Errorf("upgrade failed: current version %s, target version %s", e.agent.Version, req.TargetVersion)
+}

internal/app/ptah-agent/parse_task.go

@@ -27,6 +27,12 @@ func parseTask(taskType int, payload string) (interface{}, error) {
 		return unmarshalTask(payload, &ptahClient.UpdateCurrentNodeReq{})
 	case 8:
 		return unmarshalTask(payload, &ptahClient.DeleteServiceReq{})
+	case 9:
+		return unmarshalTask(payload, &ptahClient.DownloadAgentUpgradeReq{})
+	case 10:
+		return unmarshalTask(payload, &ptahClient.UpdateAgentSymlinkReq{})
+	case 11:
+		return unmarshalTask(payload, &ptahClient.ConfirmAgentUpgradeReq{})
 	default:
 		return nil, fmt.Errorf("parse task: unknown task type %d", taskType)
 	}

internal/app/ptah-agent/agent.go → internal/app/ptah-agent/ptah_client.go

@@ -13,16 +13,18 @@ import (
 )
 
 type Agent struct {
-	version string
+	Version string
 	ptah    *ptahClient.Client
+	rootDir string
 	docker  *dockerClient.Client
 	caddy   *caddyClient.Client
 }
 
-func New(version string, baseUrl string, ptahToken string) *Agent {
+func New(version string, baseUrl string, ptahToken string, rootDir string) *Agent {
 	return &Agent{
-		version: version,
+		Version: version,
 		ptah:    ptahClient.New(baseUrl, ptahToken),
+		rootDir: rootDir,
 		caddy:   caddyClient.New("http://127.0.0.1:2019", http.DefaultClient),
 	}
 }
@@ -48,7 +50,7 @@ func (a *Agent) sendStartedEvent(ctx context.Context) (*ptahClient.StartedRes, e
 	}
 
 	startedReq := ptahClient.StartedReq{
-		Version: a.version,
+		Version: a.Version,
 	}
 
 	startedReq.Docker.Platform.Name = info.Name
@@ -70,8 +72,12 @@ func (a *Agent) Start(ctx context.Context) error {
 	}
 
 	executor := &taskExecutor{
-		docker: a.docker,
-		caddy:  a.caddy,
+		docker:  a.docker,
+		caddy:   a.caddy,
+		rootDir: a.rootDir,
+		// TODO: use channel instead?
+		stopAgentFlag: false,
+		agent:         a,
 	}
 
 	log.Println("connected to server, poll interval", settings.Settings.PollInterval)
@@ -101,7 +107,15 @@ func (a *Agent) Start(ctx context.Context) error {
 				log.Println("can't fail task", err)
 			}
 		}
+
+		if executor.stopAgentFlag {
+			log.Println("received stop signal, shutting down gracefully")
+
+			break
+		}
 	}
+
+	return nil
 }
 
 func (a *Agent) getNextTask(ctx context.Context) (taskId int, task interface{}, err error) {

internal/app/ptah-agent/task_executor.go

@@ -11,8 +11,11 @@ import (
 )
 
 type taskExecutor struct {
-	docker *dockerClient.Client
-	caddy  *caddyClient.Client
+	agent         *Agent
+	docker        *dockerClient.Client
+	caddy         *caddyClient.Client
+	rootDir       string
+	stopAgentFlag bool
 }
 
 func (e *taskExecutor) executeTask(ctx context.Context, task interface{}) (interface{}, error) {
@@ -35,7 +38,17 @@ func (e *taskExecutor) executeTask(ctx context.Context, task interface{}) (inter
 		return e.updateCurrentNode(ctx, task.(*t.UpdateCurrentNodeReq))
 	case *t.DeleteServiceReq:
 		return e.deleteDockerService(ctx, task.(*t.DeleteServiceReq))
+	case *t.DownloadAgentUpgradeReq:
+		return e.downloadAgentUpgrade(ctx, task.(*t.DownloadAgentUpgradeReq))
+	case *t.UpdateAgentSymlinkReq:
+		return e.updateAgentSymlink(ctx, task.(*t.UpdateAgentSymlinkReq))
+	case *t.ConfirmAgentUpgradeReq:
+		return e.confirmAgentUpgrade(ctx, task.(*t.ConfirmAgentUpgradeReq))
 	default:
 		return nil, fmt.Errorf("execute task: unknown task type %T", task)
 	}
 }
+
+func (e *taskExecutor) stop() {
+	e.stopAgentFlag = true
+}

internal/pkg/ptah-client/task_types.go

@@ -105,3 +105,27 @@ type DeleteServiceReq struct {
 
 type DeleteServiceRes struct {
 }
+
+type DownloadAgentUpgradeReq struct {
+	TargetVersion string
+	DownloadUrl   string
+}
+
+type DownloadAgentUpgradeRes struct {
+	FileSize     int `json:"fileSize"`
+	DownloadTime int `json:"downloadTime"`
+}
+
+type UpdateAgentSymlinkReq struct {
+	TargetVersion string
+}
+
+type UpdateAgentSymlinkRes struct {
+}
+
+type ConfirmAgentUpgradeReq struct {
+	TargetVersion string
+}
+
+type ConfirmAgentUpgradeRes struct {
+}