perf: upload remain session replay

jumpserver · Sep 5, 2024 · 00a2a38 · 00a2a38
1 parent 108c2da
commit 00a2a38
Show file tree

Hide file tree

Showing 4 changed files with 46 additions and 9 deletions.
diff --git a/main.go b/main.go
@@ -324,9 +324,11 @@ func registerRouter(jmsService *service.JMService, tunnelService *tunnel.Guacamo
 func bootstrap(jmsService *service.JMService) {
 	replayDir := config.GlobalConfig.RecordPath
 	ftpFilePath := config.GlobalConfig.FTPFilePath
+	sessionDir := config.GlobalConfig.SessionFolderPath
 	allRemainFiles := scanRemainReplay(jmsService, replayDir)
 	go uploadRemainReplay(jmsService, allRemainFiles)
 	go uploadRemainFTPFile(jmsService, ftpFilePath)
+	go uploadRemainSessionPartReplay(jmsService, sessionDir)
 }
 
 func uploadRemainFTPFile(jmsService *service.JMService, fileStoreDir string) {
@@ -442,6 +444,30 @@ func scanRemainReplay(jmsService *service.JMService, replayDir string) map[strin
 	return allRemainFiles
 }
 
+func uploadRemainSessionPartReplay(jmsService *service.JMService, sessionDir string) {
+	sessions, err := os.ReadDir(sessionDir)
+	if err != nil {
+		logger.Errorf("Read session dir failed: %s", err)
+		return
+	}
+	terminalConf, _ := jmsService.GetTerminalConfig()
+	for _, sessionEntry := range sessions {
+		sessionId := sessionEntry.Name()
+		if !common.ValidUUIDString(sessionId) {
+			continue
+		}
+		sessionRootPath := filepath.Join(sessionDir, sessionId)
+		uploader := tunnel.PartUploader{
+			RootPath:  sessionRootPath,
+			SessionId: sessionId,
+			ApiClient: jmsService,
+			TermCfg:   &terminalConf,
+		}
+		uploader.Start()
+		logger.Infof("Upload remain session part replay %s finish", sessionId)
+	}
+}
+
 func GetStatusData(tunnelCache *tunnel.GuaTunnelCacheManager) interface{} {
 	sids := tunnelCache.RangeActiveSessionIds()
 	payload := model.HeartbeatData{

diff --git a/pkg/jms-sdk-go/model/session.go b/pkg/jms-sdk-go/model/session.go
@@ -29,16 +29,21 @@ const (
 	UnKnown  ReplayVersion = ""
 	Version2 ReplayVersion = "2"
 	Version3 ReplayVersion = "3"
+	Version5 ReplayVersion = "5"
+	Version6 ReplayVersion = "5"
 )
 
 const (
-	SuffixReplayGz = ".replay.gz"
-	SuffixCastGz   = ".cast.gz"
+	SuffixReplayGz   = ".replay.gz"
+	SuffixCastGz     = ".cast.gz"
+	SuffixPartReplay = ".part.gz"
+	SuffixReplayJson = ".replay.json"
 )
 
 var SuffixMap = map[ReplayVersion]string{
 	Version2: SuffixReplayGz,
 	Version3: SuffixCastGz,
+	Version5: SuffixPartReplay,
 }
 
 func ParseReplayVersion(gzFile string, defaultValue ReplayVersion) ReplayVersion {
@@ -47,6 +52,9 @@ func ParseReplayVersion(gzFile string, defaultValue ReplayVersion) ReplayVersion
 			return version
 		}
 	}
+	if strings.HasSuffix(gzFile, SuffixReplayJson) {
+		return Version5
+	}
 	return defaultValue
 }
 

diff --git a/pkg/tunnel/replay_part_upload.go b/pkg/tunnel/replay_part_upload.go
@@ -30,6 +30,8 @@ upload
 ├── e32248ce-2dc8-43c8-b37e-a61d5ee32176.0.part.gz
 */
 
+const ReplayType = "guacamole"
+
 type SessionReplayMeta struct {
 	model.Session
 	DateEnd    common.UTCTime `json:"date_end,omitempty"`
@@ -46,8 +48,8 @@ type PartFileMeta struct {
 type PartUploader struct {
 	SessionId string
 	RootPath  string
-	apiClient *service.JMService
-	termCfg   *model.TerminalConfig
+	ApiClient *service.JMService
+	TermCfg   *model.TerminalConfig
 
 	replayMeta SessionReplayMeta
 	partFiles  []os.DirEntry
@@ -73,7 +75,7 @@ func (p *PartUploader) preCheckSessionMeta() error {
 		endTime := GetMaxModTime(p.partFiles)
 		p.replayMeta.DateEnd = common.NewUTCTime(endTime)
 		// api finish time
-		if err1 := p.apiClient.SessionFinished(p.SessionId, p.replayMeta.DateEnd); err1 != nil {
+		if err1 := p.ApiClient.SessionFinished(p.SessionId, p.replayMeta.DateEnd); err1 != nil {
 			logger.Errorf("PartUploader %s finish session error: %v", p.SessionId, err1)
 			return err
 		}
@@ -83,6 +85,7 @@ func (p *PartUploader) preCheckSessionMeta() error {
 			logger.Errorf("PartUploader %s write meta file error: %v", p.SessionId, err1)
 		}
 	}
+	p.replayMeta.ReplayType = ReplayType
 	return nil
 }
 
@@ -183,7 +186,7 @@ func (p *PartUploader) CollectionPartFiles() {
 }
 
 func (p *PartUploader) GetStorage() storage.ReplayStorage {
-	return storage.NewReplayStorage(p.apiClient, p.termCfg.ReplayStorage)
+	return storage.NewReplayStorage(p.ApiClient, p.TermCfg.ReplayStorage)
 }
 
 const recordDirTimeFormat = "2006-01-02"
@@ -212,7 +215,7 @@ func (p *PartUploader) uploadToStorage(uploadPath string) {
 		}
 		logger.Debugf("PartUploader %s upload file %s success", p.SessionId, uploadFilePath)
 	}
-	if err = p.apiClient.FinishReply(p.SessionId); err != nil {
+	if err = p.ApiClient.FinishReply(p.SessionId); err != nil {
 		logger.Errorf("PartUploader %s finish replay error: %v", p.SessionId, err)
 		return
 	}

diff --git a/pkg/tunnel/replay_recorder.go b/pkg/tunnel/replay_recorder.go
@@ -89,8 +89,8 @@ func (r *ReplayRecorder) Stop() {
 	uploader := PartUploader{
 		RootPath:  r.RootPath,
 		SessionId: r.SessionId,
-		apiClient: r.apiClient,
-		termCfg:   r.tunnelSession.TerminalConfig,
+		ApiClient: r.apiClient,
+		TermCfg:   r.tunnelSession.TerminalConfig,
 	}
 	go uploader.Start()