nex/backend/cmd/desktop/tray.go

package main

import (
	"fmt"
	"runtime"
	"sync"
	"time"

	"nex/embedfs"

	"github.com/getlantern/systray"
	"go.uber.org/zap"
)

const defaultTrayReadyTimeout = 5 * time.Second

type trayMenuItem interface {
	Disable()
	Clicked() <-chan struct{}
}

type trayController interface {
	Run(onReady func(), onExit func())
	Quit()
	SetIcon(icon []byte)
	SetTooltip(tooltip string)
	AddMenuItem(title, tooltip string) trayMenuItem
	AddSeparator()
}

type realTrayController struct{}

func (realTrayController) Run(onReady func(), onExit func()) {
	systray.Run(onReady, onExit)
}

func (realTrayController) Quit() {
	systray.Quit()
}

func (realTrayController) SetIcon(icon []byte) {
	systray.SetIcon(icon)
}

func (realTrayController) SetTooltip(tooltip string) {
	systray.SetTooltip(tooltip)
}

func (realTrayController) AddMenuItem(title, tooltip string) trayMenuItem {
	return realTrayMenuItem{item: systray.AddMenuItem(title, tooltip)}
}

func (realTrayController) AddSeparator() {
	systray.AddSeparator()
}

type realTrayMenuItem struct {
	item *systray.MenuItem
}

func (m realTrayMenuItem) Disable() {
	m.item.Disable()
}

func (m realTrayMenuItem) Clicked() <-chan struct{} {
	return m.item.ClickedCh
}

type trayOptions struct {
	controller   trayController
	readyTimeout time.Duration
	iconLoader   func() ([]byte, error)
	openBrowser  func(string) error
	notify       func(string, string)
	logger       *zap.Logger
	fatalErrCh   <-chan error
}

func setupSystray(port int, fatalErrCh <-chan error) error {
	return runSystray(port, trayOptions{
		controller:   realTrayController{},
		readyTimeout: defaultTrayReadyTimeout,
		iconLoader:   loadTrayIcon,
		openBrowser:  openBrowser,
		notify:       showError,
		logger:       dialogLogger(),
		fatalErrCh:   fatalErrCh,
	})
}

func runSystray(port int, opts trayOptions) error {
	if opts.controller == nil {
		opts.controller = realTrayController{}
	}
	if opts.readyTimeout <= 0 {
		opts.readyTimeout = defaultTrayReadyTimeout
	}
	if opts.iconLoader == nil {
		opts.iconLoader = loadTrayIcon
	}
	if opts.openBrowser == nil {
		opts.openBrowser = openBrowser
	}
	if opts.notify == nil {
		opts.notify = showError
	}
	if opts.logger == nil {
		opts.logger = dialogLogger()
	}

	readyCh := make(chan struct{})
	doneCh := make(chan struct{})
	errCh := make(chan error, 1)
	var readyOnce sync.Once
	var errOnce sync.Once

	signalReady := func() {
		readyOnce.Do(func() { close(readyCh) })
	}
	signalError := func(err error) {
		errOnce.Do(func() { errCh <- err })
	}

	go monitorTrayStartup(port, opts, readyCh, doneCh, signalError)

	opts.controller.Run(func() {
		handleTrayReady(port, opts, signalReady, signalError)
	}, nil)
	close(doneCh)

	select {
	case err := <-errCh:
		return err
	default:
		return nil
	}
}

func monitorTrayStartup(port int, opts trayOptions, readyCh <-chan struct{}, doneCh <-chan struct{}, signalError func(error)) {
	timer := time.NewTimer(opts.readyTimeout)
	defer timer.Stop()

	ready := false
	for {
		select {
		case <-readyCh:
			ready = true
			if !timer.Stop() {
				select {
				case <-timer.C:
				default:
				}
			}
			openDesktopBrowser(port, opts)
			readyCh = nil
		case <-timer.C:
			if !ready {
				signalError(newStartupError(phaseTray, "托盘初始化超时", fmt.Errorf("托盘未在 %s 内 ready", opts.readyTimeout)))
				opts.controller.Quit()
			}
		case err := <-opts.fatalErrCh:
			if err != nil {
				signalError(newStartupError(phaseServer, startupServerErrorMessage(), err))
				opts.controller.Quit()
			}
		case <-doneCh:
			return
		}
	}
}

func handleTrayReady(port int, opts trayOptions, signalReady func(), signalError func(error)) {
	defer func() {
		if recovered := recover(); recovered != nil {
			err := fmt.Errorf("托盘初始化 panic: %v", recovered)
			signalError(newStartupError(phaseTray, "托盘菜单初始化失败", err))
			opts.controller.Quit()
		}
	}()

	icon, err := opts.iconLoader()
	if err != nil {
		signalError(newStartupError(phaseTray, "托盘图标资源无法加载", err))
		opts.controller.Quit()
		return
	}

	opts.controller.SetIcon(icon)
	opts.controller.SetTooltip(appTooltip)

	mOpen := opts.controller.AddMenuItem("打开管理界面", "在浏览器中打开")
	opts.controller.AddSeparator()
	mStatus := opts.controller.AddMenuItem("状态: 运行中", "")
	mStatus.Disable()
	mPort := opts.controller.AddMenuItem(desktopPortMenuTitle(port), "")
	mPort.Disable()
	opts.controller.AddSeparator()
	mQuit := opts.controller.AddMenuItem("退出", "停止服务并退出")

	go func() {
		for {
			select {
			case <-mOpen.Clicked():
				if err := opts.openBrowser(desktopURL(port)); err != nil {
					opts.logger.Warn("打开浏览器失败", zap.Error(err))
				}
			case <-mQuit.Clicked():
				doShutdown()
				opts.controller.Quit()
				return
			}
		}
	}()

	signalReady()
}

func openDesktopBrowser(port int, opts trayOptions) {
	url := desktopURL(port)
	if err := opts.openBrowser(url); err != nil {
		opts.logger.Warn("无法打开浏览器", zap.Error(err))
		opts.notify(appName, fmt.Sprintf("无法自动打开浏览器，请手动访问 %s", url))
	}
}

func loadTrayIcon() ([]byte, error) {
	if runtime.GOOS == "windows" {
		return embedfs.Assets.ReadFile("assets/icon.ico")
	}
	return embedfs.Assets.ReadFile("assets/icon.png")
}