myxce/src/frontend/apiserver/controllers/output.go

package controllers

import (
	"fmt"
	"strconv"
	"strings"

	"apiserver/backend"
	"apiserver/common"
	"apiserver/models"
	"apiserver/store"

	logger "apiserver/utils"

	"github.com/gin-gonic/gin"
	"go.uber.org/zap"
)

// fetchFile is the function used to fetch remote file content.
// Overrideable in tests.
var fetchFile = backend.FetchFile

const (
	// tailChunkSize is the maximum bytes fetched per request in tail mode.
	tailChunkSize = 4 * 1024 * 1024
	// incrementalChunkSize is the maximum bytes fetched in incremental mode.
	incrementalChunkSize = 4 * 1024 * 1024
)

// OutputData represents the response payload for job output.
type OutputData struct {
	Content   string `json:"content"`
	Offset    int64  `json:"offset"`
	JobStatus string `json:"job_status"`
	Done      bool   `json:"done"`
}

// GetJobOutput returns the tail or incremental content of a job's output file.
// GET /xce/v1/jobs/:jobid/output?offset=0&lines=100&type=stdout
func (j JobController) GetJobOutput(ctx *gin.Context) {
	logger.Debug("Enter GetJobOutput() function from ", zap.String("client_ip", ctx.ClientIP()))

	token := ctx.GetHeader("Authorization")
	submitUser, err := GetUsernameFromToken(token)
	if err != nil {
		ReturnError(ctx, common.ERROR_CODE, err.Error(), nil)
		return
	}

	if ctx.IsAborted() {
		ReturnError(ctx, common.ERROR_AUTH_CODE, common.ERROR_ACCESS_FORBIDDEN_MSG, nil)
		return
	}

	jobidStr := ctx.Param("jobids")
	offsetStr := ctx.DefaultQuery("offset", "0")
	linesStr := ctx.DefaultQuery("lines", "100")
	streamType := ctx.DefaultQuery("type", "stdout")

	offset, err := strconv.ParseInt(offsetStr, 10, 64)
	if err != nil {
		ReturnError(ctx, common.INVALID_PARAMS_CODE, "invalid offset: "+err.Error(), nil)
		return
	}

	lines, err := strconv.Atoi(linesStr)
	if err != nil || lines <= 0 {
		ReturnError(ctx, common.INVALID_PARAMS_CODE, "invalid lines: must be a positive integer", nil)
		return
	}

	if streamType != "stdout" && streamType != "stderr" {
		ReturnError(ctx, common.INVALID_PARAMS_CODE, "invalid type: must be 'stdout' or 'stderr'", nil)
		return
	}

	// Parse job ID
	jobInfo, err := parseJobIDParam(jobidStr)
	if err != nil {
		ReturnError(ctx, common.INVALID_PARAMS_CODE, common.INVALID_PARAMS_MSG+": "+err.Error(), nil)
		return
	}

	// Look up job in store
	job, ok := store.Jobs.Get(uint64(jobInfo.Unique_JobID))
	if !ok {
		ReturnError(ctx, common.ERROR_NOT_FOUND_CODE, common.ERROR_JOB_NOT_FOUND_MSG, nil)
		return
	}

	// Check ownership: the owner may always read; anyone else needs a
	// job:view grant covering the job (bug-1934)
	if job.SubmitUser != submitUser {
		scope, perr := userJobScope(submitUser, "job:view")
		if perr != nil || !scope.covers(job.SubmitUser) {
			ReturnError(ctx, common.ERROR_AUTH_CODE, common.ERROR_ACCESS_FORBIDDEN_MSG, nil)
			return
		}
	}

	// Resolve output file path
	var filePath string
	if streamType == "stderr" {
		filePath = resolveOutputPath(job.ErrFile, job)
	} else {
		filePath = resolveOutputPath(job.OutFile, job)
	}

	if filePath == "" {
		ReturnError(ctx, common.ERROR_READ_FILE_CODE, fmt.Sprintf("job <%s>: no output file path could be determined for %s", backend.FormatJobID(job.Unique_JobID), streamType), nil)
		return
	}

	var content string
	var newOffset int64

	if offset == 0 {
		content, newOffset, err = fetchFileTail(uint64(job.Unique_JobID), filePath, lines)
	} else {
		content, newOffset, err = fetchFileFromOffset(uint64(job.Unique_JobID), filePath, offset)
	}

	if err != nil {
		ReturnError(ctx, common.ERROR_READ_FILE_CODE, fmt.Sprintf("job <%s> (path=%s): failed to read output: %s", backend.FormatJobID(job.Unique_JobID), filePath, err.Error()), nil)
		return
	}

	jobStatusStr := models.JobStateStr(job.JobStatus)
	done := jobStatusStr != "RUN" && jobStatusStr != "PSUSP" && jobStatusStr != "USUSP" && jobStatusStr != "SSUSP"

	ReturnSuccess(ctx, common.SUCCESS_CODE, common.SUCCESS_MSG, OutputData{
		Content:   content,
		Offset:    newOffset,
		JobStatus: jobStatusStr,
		Done:      done,
	})
}

// fetchFileTail fetches the last n lines of a file via the backend.
// It uses a two-pass approach for files larger than tailChunkSize.
func fetchFileTail(jobid uint64, path string, n int) (string, int64, error) {
	// First pass: fetch from offset 0 to get file size and content for small files.
	result, err := fetchFile(jobid, path, 0, tailChunkSize)
	if err != nil {
		return "", 0, err
	}

	fileSize := int64(result.FileSize)
	if fileSize == 0 {
		return "", 0, nil
	}

	data := result.Content

	// If the file is larger than our chunk, re-fetch the tail portion.
	if result.FileSize > tailChunkSize {
		endOffset := result.FileSize - tailChunkSize
		result, err = fetchFile(jobid, path, endOffset, tailChunkSize)
		if err != nil {
			return "", 0, err
		}
		data = result.Content
	}

	return lastNLines(data, n), fileSize, nil
}

// fetchFileFromOffset fetches content from the given byte offset via the backend.
func fetchFileFromOffset(jobid uint64, path string, offset int64) (string, int64, error) {
	result, err := fetchFile(jobid, path, uint64(offset), incrementalChunkSize)
	if err != nil {
		return "", 0, err
	}

	fileSize := int64(result.FileSize)
	if int64(offset) >= fileSize {
		return "", fileSize, nil
	}

	return string(result.Content), fileSize, nil
}

// lastNLines returns the last n lines from data as a string.
func lastNLines(data []byte, n int) string {
	if len(data) == 0 {
		return ""
	}

	lineCount := 0
	start := 0
	end := len(data) - 1

	for i := end; i >= 0; i-- {
		if data[i] == '\n' {
			lineCount++
			if lineCount == n+1 {
				start = i + 1
				break
			}
		}
	}

	return string(data[start:])
}

// resolveOutputPath replaces placeholders in the output file path.
// When path is empty, returns the worker's default output path for task 0,
// matching uv_proc.c make_job_output_path() behavior.
func resolveOutputPath(path string, job models.JobInfo) string {
	if path == "" {
		return backend.WorkerDefaultOutputPath(job.Unique_JobID, 0)
	}
	formattedID := backend.FormatJobID(job.Unique_JobID)
	path = strings.ReplaceAll(path, "%J", formattedID)
	path = strings.ReplaceAll(path, "%I", strconv.Itoa(job.ArrayIdx))
	path = strings.ReplaceAll(path, "%U", job.SubmitUser)
	if job.CWD != "" {
		path = strings.ReplaceAll(path, "%C", job.CWD)
	}
	return path
}

// parseJobIDParam parses the job peek names and returns the corresponding
// JobInfo: a job number (an ordinary job, or an element by its own number,
// XCE_JOBID) names that job exactly; <array-id>[k] names element k. A range
// or an invalid selection is an error -- peek shows one job.
func parseJobIDParam(jobidStr string) (models.JobInfo, error) {
	raw := strings.TrimSpace(jobidStr)
	if n, err := strconv.ParseUint(raw, 10, 64); err == nil {
		return models.JobInfo{Unique_JobID: int64(n), ArrayIdx: backend.ExtractArrayIdx(n)}, nil
	}

	sel, selErr := ParseIdentitySelection(raw)
	if selErr != nil {
		return models.JobInfo{}, selErr
	}
	switch sel.Form {
	case SelSingle:
		return models.JobInfo{
			Unique_JobID: int64(backend.MakeElementID(sel.Key, sel.Lo)),
			ArrayIdx:     sel.Lo,
		}, nil
	case SelRange:
		return models.JobInfo{}, fmt.Errorf("'%s' names several jobs; peek one element as <array-id>[k]", raw)
	default:
		return models.JobInfo{Unique_JobID: int64(sel.Key)}, nil
	}
}