github.com/nikandfor/tlog@v0.21.5-0.20231108111739-3ef89426a96d/agent/query.go (about)

     1  package agent
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"io"
     7  
     8  	"github.com/nikandfor/tlog"
     9  )
    10  
    11  func (a *Agent) Query(ctx context.Context, w io.Writer, q string) (err error) {
    12  	tr := tlog.SpawnFromContextOrStart(ctx, "agent_query", "query", q)
    13  	defer tr.Finish("err", &err)
    14  
    15  	ctx = tlog.ContextWithSpan(ctx, tr)
    16  
    17  	id := a.registerSub(w)
    18  	defer a.unregisterSub(id)
    19  
    20  	<-ctx.Done()
    21  	err = ctx.Err()
    22  
    23  	if errors.Is(err, context.Canceled) {
    24  		err = nil
    25  	}
    26  
    27  	return
    28  }
    29  
    30  func (a *Agent) registerSub(w io.Writer) int64 {
    31  	defer a.Unlock()
    32  	a.Lock()
    33  
    34  	a.subid++
    35  	id := a.subid
    36  
    37  	s := &sub{
    38  		id:     id,
    39  		Writer: w,
    40  	}
    41  
    42  	a.subs[id] = s
    43  
    44  	return id
    45  }
    46  
    47  func (a *Agent) unregisterSub(id int64) {
    48  	defer a.Unlock()
    49  	a.Lock()
    50  
    51  	delete(a.subs, id)
    52  }