129 lines
3.3 KiB
Go
129 lines
3.3 KiB
Go
package transfer
|
|
|
|
import (
|
|
"context"
|
|
"time"
|
|
|
|
"github.com/apex/log"
|
|
"github.com/mitchellh/colorstring"
|
|
|
|
"github.com/pterodactyl/wings/server"
|
|
"github.com/pterodactyl/wings/system"
|
|
)
|
|
|
|
// Status represents the current status of a transfer.
|
|
type Status string
|
|
|
|
// String satisfies the fmt.Stringer interface.
|
|
func (s Status) String() string {
|
|
return string(s)
|
|
}
|
|
|
|
const (
|
|
// StatusPending is the status of a transfer when it is first created.
|
|
StatusPending Status = "pending"
|
|
// StatusProcessing is the status of a transfer when it is currently in
|
|
// progress, such as when the archive is being streamed to the target node.
|
|
StatusProcessing Status = "processing"
|
|
|
|
// StatusCancelling is the status of a transfer when it is in the process of
|
|
// being cancelled.
|
|
StatusCancelling Status = "cancelling"
|
|
|
|
// StatusCancelled is the final status of a transfer when it has been
|
|
// cancelled.
|
|
StatusCancelled Status = "cancelled"
|
|
// StatusFailed is the final status of a transfer when it has failed.
|
|
StatusFailed Status = "failed"
|
|
// StatusCompleted is the final status of a transfer when it has completed.
|
|
StatusCompleted Status = "completed"
|
|
)
|
|
|
|
// Transfer represents a transfer of a server from one node to another.
|
|
type Transfer struct {
|
|
// ctx is the context for the transfer.
|
|
ctx context.Context
|
|
// cancel is used to cancel all ongoing transfer operations for the server.
|
|
cancel *context.CancelFunc
|
|
|
|
// Server associated with the transfer.
|
|
Server *server.Server
|
|
// status of the transfer.
|
|
status *system.Atomic[Status]
|
|
|
|
// archive is the archive that is being created for the transfer.
|
|
archive *Archive
|
|
}
|
|
|
|
// New returns a new transfer instance for the given server.
|
|
func New(ctx context.Context, s *server.Server) *Transfer {
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
|
|
return &Transfer{
|
|
ctx: ctx,
|
|
cancel: &cancel,
|
|
|
|
Server: s,
|
|
status: system.NewAtomic(StatusPending),
|
|
}
|
|
}
|
|
|
|
// Context returns the context for the transfer.
|
|
func (t *Transfer) Context() context.Context {
|
|
return t.ctx
|
|
}
|
|
|
|
// Cancel cancels the transfer.
|
|
func (t *Transfer) Cancel() {
|
|
status := t.Status()
|
|
if status == StatusCancelling ||
|
|
status == StatusCancelled ||
|
|
status == StatusCompleted ||
|
|
status == StatusFailed {
|
|
return
|
|
}
|
|
|
|
if t.cancel == nil {
|
|
return
|
|
}
|
|
|
|
t.SetStatus(StatusCancelling)
|
|
(*t.cancel)()
|
|
}
|
|
|
|
// Status returns the current status of the transfer.
|
|
func (t *Transfer) Status() Status {
|
|
return t.status.Load()
|
|
}
|
|
|
|
// SetStatus sets the status of the transfer.
|
|
func (t *Transfer) SetStatus(s Status) {
|
|
// TODO: prevent certain status changes from happening.
|
|
// If we are cancelling, then we can't go back to processing.
|
|
t.status.Store(s)
|
|
|
|
t.Server.Events().Publish(server.TransferStatusEvent, s)
|
|
}
|
|
|
|
// SendMessage sends a message to the server's console.
|
|
func (t *Transfer) SendMessage(v string) {
|
|
t.Server.Events().Publish(
|
|
server.TransferLogsEvent,
|
|
colorstring.Color("[yellow][bold]"+time.Now().Format(time.RFC1123)+" [Transfer System] [Source Node]:[default] "+v),
|
|
)
|
|
}
|
|
|
|
// Error logs an error that occurred on the source node.
|
|
func (t *Transfer) Error(err error, v string) {
|
|
t.Log().WithError(err).Error(v)
|
|
t.SendMessage(v)
|
|
}
|
|
|
|
// Log returns a logger for the transfer.
|
|
func (t *Transfer) Log() *log.Entry {
|
|
if t.Server == nil {
|
|
return log.WithField("subsystem", "transfer")
|
|
}
|
|
return t.Server.Log().WithField("subsystem", "transfer")
|
|
}
|