uopi/pkg/pva/client.go

package pva

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"log/slog"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/uopi/gopva/pvdata"
)

// DefaultTCPPort is the standard PVA TCP port for channel connections.
const DefaultTCPPort = 5075

// DefaultUDPPort is the PVA broadcast/search port (distinct from TCP).
const DefaultUDPPort = 5076

// Client is a PVA client that manages connections to one or more servers.
// It performs UDP search to locate PVs, then opens TCP connections on demand.
type Client struct {
	mu      sync.Mutex
	conns   map[string]*conn    // server addr → connection
	pending map[string][]waiter // pvName → waiters for address

	searchSeq atomic.Uint32
	udp       *net.UDPConn

	ctx    context.Context
	cancel context.CancelFunc
}

type waiter struct {
	pvName string
	cb     func(addr string, err error)
}

// NewClient creates a PVA client.  Call Close when done.
func NewClient() (*Client, error) {
	ctx, cancel := context.WithCancel(context.Background())
	cl := &Client{
		conns:   make(map[string]*conn),
		pending: make(map[string][]waiter),
		ctx:     ctx,
		cancel:  cancel,
	}
	if err := cl.startUDP(); err != nil {
		cancel()
		return nil, err
	}
	return cl, nil
}

// Close shuts down the client and all connections.
func (cl *Client) Close() {
	cl.cancel()
	if cl.udp != nil {
		cl.udp.Close()
	}
	cl.mu.Lock()
	defer cl.mu.Unlock()
	for _, c := range cl.conns {
		c.close()
	}
}

// Get issues a one-shot GET for pvName and returns the decoded structure.
func (cl *Client) Get(ctx context.Context, pvName string) (pvdata.StructValue, error) {
	ch := make(chan struct {
		v   pvdata.StructValue
		err error
	}, 1)
	cl.withConn(ctx, pvName, func(c *conn, err error) {
		if err != nil {
			ch <- struct {
				v   pvdata.StructValue
				err error
			}{err: err}
			return
		}
		c.get(pvName, func(v pvdata.StructValue, e error) {
			ch <- struct {
				v   pvdata.StructValue
				err error
			}{v, e}
		})
	})
	select {
	case res := <-ch:
		return res.v, res.err
	case <-ctx.Done():
		return pvdata.StructValue{}, ctx.Err()
	}
}

// Monitor subscribes to pvName and returns a channel that receives updates.
// The channel is closed when ctx is cancelled or a fatal error occurs.
func (cl *Client) Monitor(ctx context.Context, pvName string) <-chan MonitorEvent {
	ch := make(chan MonitorEvent, 16)
	cl.withConn(ctx, pvName, func(c *conn, err error) {
		if err != nil {
			ch <- MonitorEvent{Err: err}
			close(ch)
			return
		}
		c.subscribe(ctx, pvName, ch)
	})
	return ch
}

// ---- Internal --------------------------------------------------------

// withConn locates (or creates) a connection to the server hosting pvName
// and calls cb on it.  cb may be called from a goroutine.
func (cl *Client) withConn(ctx context.Context, pvName string, cb func(*conn, error)) {
	cl.search(ctx, pvName, func(addr string, err error) {
		if err != nil {
			cb(nil, err)
			return
		}
		cl.mu.Lock()
		c, ok := cl.conns[addr]
		cl.mu.Unlock()
		if ok {
			cb(c, nil)
			return
		}
		// open new TCP connection
		tc, err := dial(ctx, addr)
		if err != nil {
			cb(nil, fmt.Errorf("pva: connect %s: %w", addr, err))
			return
		}
		cl.mu.Lock()
		cl.conns[addr] = tc
		cl.mu.Unlock()
		cb(tc, nil)
	})
}

// ---- UDP search -------------------------------------------------------

func (cl *Client) startUDP() error {
	conn, err := net.ListenUDP("udp4", &net.UDPAddr{Port: 0})
	if err != nil {
		return fmt.Errorf("pva: UDP listen: %w", err)
	}
	cl.udp = conn
	go cl.udpReadLoop()
	return nil
}

// search sends a PVA search request and calls cb when a response arrives.
// Currently does broadcast + localhost; production code would also check
// EPICS_PVA_ADDR_LIST environment variable.
func (cl *Client) search(ctx context.Context, pvName string, cb func(string, error)) {
	seq := cl.searchSeq.Add(1)

	cl.mu.Lock()
	cl.pending[pvName] = append(cl.pending[pvName], waiter{pvName: pvName, cb: cb})
	cl.mu.Unlock()

	go func() {
		for attempt := 0; attempt < 10; attempt++ {
			select {
			case <-ctx.Done():
				cl.mu.Lock()
				cl.removeWaiter(pvName, cb)
				cl.mu.Unlock()
				cb("", ctx.Err())
				return
			default:
			}
			cl.sendSearchRequest(seq, pvName)
			// exponential backoff: 100ms, 200ms, 400ms … up to 2s
			delay := time.Duration(100<<min(attempt, 4)) * time.Millisecond
			select {
			case <-time.After(delay):
			case <-ctx.Done():
				cl.mu.Lock()
				cl.removeWaiter(pvName, cb)
				cl.mu.Unlock()
				cb("", ctx.Err())
				return
			}
		}
		// give up
		cl.mu.Lock()
		cl.removeWaiter(pvName, cb)
		cl.mu.Unlock()
		cb("", fmt.Errorf("pva: search timeout for %q", pvName))
	}()
}

func (cl *Client) removeWaiter(pvName string, cb func(string, error)) {
	ws := cl.pending[pvName]
	for i, w := range ws {
		// compare by pointer (closure identity via fmt trick)
		if fmt.Sprintf("%p", w.cb) == fmt.Sprintf("%p", cb) {
			cl.pending[pvName] = append(ws[:i], ws[i+1:]...)
			return
		}
	}
}

// sendSearchRequest broadcasts a PVA SEARCH request.
func (cl *Client) sendSearchRequest(seq uint32, pvName string) {
	// SEARCH payload (spec §6.3):
	// searchSeqID(4) + flags(1) + reserved(3) + responseAddress(16) + responsePort(2)
	// + transportCount(1) + transports[](1=TCP) + pvCount(2) + pvs[](channelID(4)+name(str))
	var buf bytes.Buffer
	binary.Write(&buf, binary.LittleEndian, seq) // searchSeqID
	buf.WriteByte(0x81)                          // flags: unicast=0, reply=1, mustReply=1
	buf.Write([]byte{0, 0, 0})                   // reserved
	// responseAddress: 16 bytes (IPv4-mapped IPv6: ::ffff:0.0.0.0 = no specific address)
	buf.Write(make([]byte, 12))
	buf.Write([]byte{0, 0, 0, 0}) // 0.0.0.0 → server will reply to source addr

	// responsePort: our UDP listening port
	laddr := cl.udp.LocalAddr().(*net.UDPAddr)
	binary.Write(&buf, binary.LittleEndian, uint16(laddr.Port))

	buf.WriteByte(1)    // transportCount = 1
	buf.WriteByte(0x01) // transport[0]: TCP

	binary.Write(&buf, binary.LittleEndian, uint16(1))   // pvCount = 1
	binary.Write(&buf, binary.LittleEndian, uint32(seq)) // channelID (reuse seq)
	pvdata.WriteString(&buf, pvName)

	msg := BuildMessage(CmdSearchRequest, flagApp, buf.Bytes())

	// Send to localhost and broadcast on the PVA UDP search port (5076).
	targets := []string{
		fmt.Sprintf("127.0.0.1:%d", DefaultUDPPort),
		fmt.Sprintf("255.255.255.255:%d", DefaultUDPPort),
	}
	// Also check EPICS_PVA_ADDR_LIST if set (common in labs).
	// Entries without an explicit port default to DefaultTCPPort for TCP
	// connections; for the search broadcast we still use the UDP port.
	for _, addr := range addrsFromEnv() {
		targets = append(targets, fmt.Sprintf("%s:%d", addr, DefaultUDPPort))
	}
	for _, addr := range targets {
		dst, err := net.ResolveUDPAddr("udp4", addr)
		if err != nil {
			continue
		}
		if _, err := cl.udp.WriteTo(msg, dst); err != nil {
			slog.Debug("pva search send", "dst", addr, "err", err)
		}
	}
}

// udpReadLoop reads SEARCH_RESPONSE (and BEACON) UDP messages.
func (cl *Client) udpReadLoop() {
	buf := make([]byte, 65536)
	for {
		n, src, err := cl.udp.ReadFromUDP(buf)
		if err != nil {
			return
		}
		cl.handleUDP(buf[:n], src)
	}
}

func (cl *Client) handleUDP(data []byte, src *net.UDPAddr) {
	if len(data) < headerSize {
		return
	}
	r := bytes.NewReader(data)
	hdr, err := ReadHeader(r)
	if err != nil || hdr.isControl() {
		return
	}
	if hdr.Command != CmdSearchResponse {
		return
	}

	payload := make([]byte, hdr.Size)
	if _, err := r.Read(payload); err != nil {
		return
	}
	pr := bytes.NewReader(payload)

	// SEARCH_RESPONSE:
	// guid(12) + searchSeqID(4) + serverAddress(16) + serverPort(2)
	// + protocol(str) + found(1) + pvCount(2) + pvIDs[](4)
	var guid [12]byte
	pr.Read(guid[:])
	var seq uint32
	binary.Read(pr, binary.LittleEndian, &seq)
	var addr [16]byte
	pr.Read(addr[:])
	var port uint16
	binary.Read(pr, binary.LittleEndian, &port)
	proto, _ := readPVAString(pr, binary.LittleEndian)
	if proto != "tcp" {
		return
	}
	found, _ := pr.ReadByte()
	if found == 0 {
		return
	}
	var pvCount uint16
	binary.Read(pr, binary.LittleEndian, &pvCount)
	// pvIDs — we match by searching for waiters by name (seq used as channelID)

	// Determine server address: use src if serverAddress is 0.0.0.0
	ip := net.IP(addr[12:16]) // last 4 bytes = IPv4
	if ip.Equal(net.IPv4zero) {
		ip = src.IP
	}
	serverAddr := fmt.Sprintf("%s:%d", ip.String(), port)

	// Notify all pending waiters (simple: we got a response so all PVs are on this server)
	cl.mu.Lock()
	var toNotify []func(string, error)
	for pvName, ws := range cl.pending {
		for _, w := range ws {
			toNotify = append(toNotify, w.cb)
		}
		delete(cl.pending, pvName)
	}
	cl.mu.Unlock()

	for _, cb := range toNotify {
		go cb(serverAddr, nil)
	}
}

func min(a, b int) int {
	if a < b {
		return a
	}
	return b
}