pubsub

refactor: remote client writer that implement io.Writer rss
created pr with 31.1 on 2024-11-12T20:57:56Z · by c8ef7d19
cmds
checkout latest patchset: 
ssh pr.pico.sh print 31 | git am -3
checkout any patchset in a patch request: 
ssh pr.pico.sh print 31.[rev] | git am -3
add changes to patch request: 
git format-patch main --stdout | ssh pr.pico.sh pr add 31

Patchset 31.1 on 2024-11-12T20:57:56Z · commit de933be

→ refactor: remote client writer that implement io.Writer
Eric Bower 2024-11-12T20:57:12Z
+6 -176 log/log.go #
......@@ -3,15 +3,12 @@ package log
33 import (
44 "context"
55 "errors"
6- "fmt"
76 "io"
87 "log/slog"
98 "slices"
109 "sync"
11- "time"
1210
1311 "github.com/picosh/pubsub"
14- "golang.org/x/crypto/ssh"
1512 )
1613
1714 type MultiHandler struct {
......@@ -19,6 +16,8 @@ type MultiHandler struct {
1916 mu sync.Mutex
2017 }
2118
19+var _ slog.Handler = (*MultiHandler)(nil)
20+
2221 func (m *MultiHandler) Enabled(ctx context.Context, l slog.Level) bool {
2322 m.mu.Lock()
2423 defer m.mu.Unlock()
......@@ -80,181 +79,15 @@ func (m *MultiHandler) WithGroup(name string) slog.Handler {
8079 }
8180 }
8281
83-type PubSubLogWriter struct {
84- SSHClient *ssh.Client
85- Session *ssh.Session
86- StdinPipe io.WriteCloser
87- Done chan struct{}
88- Messages chan []byte
89- Timeout time.Duration
90- BufferSize int
91- closeOnce sync.Once
92- closeMessageOnce sync.Once
93- startOnce sync.Once
94- connecMu sync.Mutex
95- ConnectionInfo *pubsub.RemoteClientInfo
96-}
97-
98-func (c *PubSubLogWriter) Close() error {
99- c.connecMu.Lock()
100- defer c.connecMu.Unlock()
101-
102- if c.Done != nil {
103- c.closeOnce.Do(func() {
104- close(c.Done)
105- })
106- }
107-
108- if c.Messages != nil {
109- c.closeMessageOnce.Do(func() {
110- close(c.Messages)
111- })
112- }
113-
114- var errs []error
115-
116- if c.StdinPipe != nil {
117- errs = append(errs, c.StdinPipe.Close())
118- }
119-
120- if c.Session != nil {
121- errs = append(errs, c.Session.Close())
122- }
123-
124- if c.SSHClient != nil {
125- errs = append(errs, c.SSHClient.Close())
126- }
127-
128- return errors.Join(errs...)
129-}
130-
131-func (c *PubSubLogWriter) Open() error {
132- c.Close()
133-
134- c.connecMu.Lock()
135-
136- c.Done = make(chan struct{})
137- c.Messages = make(chan []byte, c.BufferSize)
138-
139- sshClient, err := pubsub.CreateRemoteClient(c.ConnectionInfo)
140- if err != nil {
141- c.connecMu.Unlock()
142- return err
143- }
144-
145- session, err := sshClient.NewSession()
146- if err != nil {
147- c.connecMu.Unlock()
148- return err
149- }
150-
151- stdinPipe, err := session.StdinPipe()
152- if err != nil {
153- c.connecMu.Unlock()
154- return err
155- }
156-
157- err = session.Start("pub log-drain -b=false")
158- if err != nil {
159- c.connecMu.Unlock()
160- return err
161- }
162-
163- c.SSHClient = sshClient
164- c.Session = session
165- c.StdinPipe = stdinPipe
166-
167- c.closeOnce = sync.Once{}
168- c.startOnce = sync.Once{}
169-
170- c.connecMu.Unlock()
171-
172- c.Start()
173-
174- return nil
175-}
176-
177-func (c *PubSubLogWriter) Start() {
178- c.startOnce.Do(func() {
179- go func() {
180- defer c.Reconnect()
181-
182- for {
183- select {
184- case data, ok := <-c.Messages:
185- _, err := c.StdinPipe.Write(data)
186- if !ok || err != nil {
187- slog.Error("received error on write, reopening logger", "error", err)
188- return
189- }
190- case <-c.Done:
191- return
192- }
193- }
194- }()
195- })
196-}
197-
198-func (c *PubSubLogWriter) Write(data []byte) (int, error) {
199- var (
200- n int
201- err error
202- )
203-
204- ok := c.connecMu.TryLock()
205-
206- if !ok {
207- return n, fmt.Errorf("unable to acquire lock to write")
208- }
209-
210- defer c.connecMu.Unlock()
211-
212- if c.Messages == nil || c.Done == nil {
213- return n, fmt.Errorf("logger not viable")
214- }
215-
216- select {
217- case c.Messages <- slices.Clone(data):
218- n = len(data)
219- case <-time.After(c.Timeout):
220- err = fmt.Errorf("unable to send data within timeout")
221- case <-c.Done:
222- break
223- }
224-
225- return n, err
226-}
227-
228-func (c *PubSubLogWriter) Reconnect() {
229- go func() {
230- for {
231- err := c.Open()
232- if err != nil {
233- slog.Error("unable to open send logger. retrying in 10 seconds", "error", err)
234- } else {
235- return
236- }
237-
238- <-time.After(10 * time.Second)
239- }
240- }()
241-}
242-
243-func SendLogRegister(logger *slog.Logger, connectionInfo *pubsub.RemoteClientInfo, buffer int) (*slog.Logger, error) {
82+func SendLogRegister(logger *slog.Logger, info *pubsub.RemoteClientInfo, buffer int) (*slog.Logger, error) {
24483 if buffer < 0 {
24584 buffer = 0
24685 }
24786
248- currentHandler := logger.Handler()
249-
250- logWriter := &PubSubLogWriter{
251- Timeout: 10 * time.Millisecond,
252- BufferSize: buffer,
253- ConnectionInfo: connectionInfo,
254- }
255-
256- logWriter.Reconnect()
87+ logWriter := pubsub.NewRemoteClientWriter(info, logger, buffer)
88+ go logWriter.KeepAlive("pub log-drain -b=false")
25789
90+ currentHandler := logger.Handler()
25891 return slog.New(
25992 &MultiHandler{
26093 Handlers: []slog.Handler{
......@@ -268,9 +101,6 @@ func SendLogRegister(logger *slog.Logger, connectionInfo *pubsub.RemoteClientInf
268101 ), nil
269102 }
270103
271-var _ io.Writer = (*PubSubLogWriter)(nil)
272-var _ slog.Handler = (*MultiHandler)(nil)
273-
274104 func ConnectToLogs(ctx context.Context, connectionInfo *pubsub.RemoteClientInfo) (io.Reader, error) {
275105 return pubsub.RemoteSub("sub log-drain -k", ctx, connectionInfo)
276106 }
+174 -1 remote_client.go #
......@@ -2,16 +2,189 @@ package pubsub
22
33 import (
44 "context"
5+ "errors"
56 "fmt"
67 "io"
8+ "log/slog"
79 "net"
810 "os"
911 "path/filepath"
12+ "slices"
1013 "strings"
14+ "sync"
15+ "time"
1116
1217 "golang.org/x/crypto/ssh"
1318 )
1419
20+type RemoteClientWriter struct {
21+ SSHClient *ssh.Client
22+ Session *ssh.Session
23+ StdinPipe io.WriteCloser
24+ Done chan struct{}
25+ Messages chan []byte
26+ Timeout time.Duration
27+ BufferSize int
28+ closeOnce sync.Once
29+ closeMessageOnce sync.Once
30+ startOnce sync.Once
31+ connecMu sync.Mutex
32+ Info *RemoteClientInfo
33+ Logger *slog.Logger
34+}
35+
36+var _ io.Writer = (*RemoteClientWriter)(nil)
37+
38+func NewRemoteClientWriter(info *RemoteClientInfo, logger *slog.Logger, buffer int) *RemoteClientWriter {
39+ return &RemoteClientWriter{
40+ Timeout: 10 * time.Millisecond,
41+ Info: info,
42+ BufferSize: buffer,
43+ Logger: logger,
44+ }
45+}
46+
47+func (c *RemoteClientWriter) Close() error {
48+ c.connecMu.Lock()
49+ defer c.connecMu.Unlock()
50+
51+ if c.Done != nil {
52+ c.closeOnce.Do(func() {
53+ close(c.Done)
54+ })
55+ }
56+
57+ if c.Messages != nil {
58+ c.closeMessageOnce.Do(func() {
59+ close(c.Messages)
60+ })
61+ }
62+
63+ var errs []error
64+
65+ if c.StdinPipe != nil {
66+ errs = append(errs, c.StdinPipe.Close())
67+ }
68+
69+ if c.Session != nil {
70+ errs = append(errs, c.Session.Close())
71+ }
72+
73+ if c.SSHClient != nil {
74+ errs = append(errs, c.SSHClient.Close())
75+ }
76+
77+ return errors.Join(errs...)
78+}
79+
80+func (c *RemoteClientWriter) Open(cmd string) error {
81+ c.Close()
82+
83+ c.connecMu.Lock()
84+
85+ c.Done = make(chan struct{})
86+ c.Messages = make(chan []byte, c.BufferSize)
87+
88+ sshClient, err := CreateRemoteClient(c.Info)
89+ if err != nil {
90+ c.connecMu.Unlock()
91+ return err
92+ }
93+
94+ session, err := sshClient.NewSession()
95+ if err != nil {
96+ c.connecMu.Unlock()
97+ return err
98+ }
99+
100+ stdinPipe, err := session.StdinPipe()
101+ if err != nil {
102+ c.connecMu.Unlock()
103+ return err
104+ }
105+
106+ err = session.Start(cmd)
107+ if err != nil {
108+ c.connecMu.Unlock()
109+ return err
110+ }
111+
112+ c.SSHClient = sshClient
113+ c.Session = session
114+ c.StdinPipe = stdinPipe
115+
116+ c.closeOnce = sync.Once{}
117+ c.startOnce = sync.Once{}
118+
119+ c.connecMu.Unlock()
120+
121+ c.start()
122+
123+ return nil
124+}
125+
126+func (c *RemoteClientWriter) start() {
127+ c.startOnce.Do(func() {
128+ go func() {
129+ for {
130+ select {
131+ case data, ok := <-c.Messages:
132+ _, err := c.StdinPipe.Write(data)
133+ if !ok || err != nil {
134+ c.Logger.Error("received error on write, reopening conn", "error", err)
135+ return
136+ }
137+ case <-c.Done:
138+ return
139+ }
140+ }
141+ }()
142+ })
143+}
144+
145+func (c *RemoteClientWriter) Write(data []byte) (int, error) {
146+ var (
147+ n int
148+ err error
149+ )
150+
151+ ok := c.connecMu.TryLock()
152+
153+ if !ok {
154+ return n, fmt.Errorf("unable to acquire lock to write")
155+ }
156+
157+ defer c.connecMu.Unlock()
158+
159+ if c.Messages == nil || c.Done == nil {
160+ return n, fmt.Errorf("conn not viable")
161+ }
162+
163+ select {
164+ case c.Messages <- slices.Clone(data):
165+ n = len(data)
166+ case <-time.After(c.Timeout):
167+ err = fmt.Errorf("unable to send data within timeout")
168+ case <-c.Done:
169+ break
170+ }
171+
172+ return n, err
173+}
174+
175+func (c *RemoteClientWriter) KeepAlive(cmd string) {
176+ for {
177+ err := c.Open(cmd)
178+ if err != nil {
179+ c.Logger.Error("unable to open send to ssh conn. retrying in 10 seconds", "error", err)
180+ } else {
181+ return
182+ }
183+
184+ <-time.After(10 * time.Second)
185+ }
186+}
187+
15188 type RemoteClientInfo struct {
16189 RemoteHost string
17190 KeyLocation string
......@@ -22,7 +195,7 @@ type RemoteClientInfo struct {
22195
23196 func CreateRemoteClient(info *RemoteClientInfo) (*ssh.Client, error) {
24197 if info == nil {
25- return nil, fmt.Errorf("connection info is invalid")
198+ return nil, fmt.Errorf("conn info is invalid")
26199 }
27200
28201 if !strings.Contains(info.RemoteHost, ":") {
Back to top