add shutdown

Author: emranbm

errors.go

@@ -3,3 +3,4 @@ package golangfuse
 import "errors"
 
 var AlreadyStartedErr = errors.New("already started")
+var AlreadyShutdownErr = errors.New("already shutdown")

event_buffer.go

@@ -26,7 +26,7 @@ func (i *eventBuffer) Start(ctx context.Context, period time.Duration) {
 	for {
 		select {
 		case <-ticker.C:
-			i.flush(ctx)
+			i.Flush(ctx)
 		case <-ctx.Done():
 			return
 		}
@@ -39,7 +39,7 @@ func (i *eventBuffer) Add(event IngestionEvent) {
 	i.bufferedEvents = append(i.bufferedEvents, event)
 }
 
-func (i *eventBuffer) flush(ctx context.Context) {
+func (i *eventBuffer) Flush(ctx context.Context) {
 	i.mu.Lock()
 	items := i.bufferedEvents
 	i.bufferedEvents = nil

langfuse.go

@@ -14,6 +14,9 @@ import (
 
 type Langfuse interface {
 	StartSendingEvents(ctx context.Context, period time.Duration) error
+	// Shutdown stops the event buffer and cleans up resources.
+	// It also ensures that any pending events are sent before shutdown.
+	Shutdown(ctx context.Context) error
 	Trace(input, output any, options ...TraceOption)
 	GetSystemPromptTemplate(ctx context.Context, promptName string) (string, error)
 }
@@ -22,6 +25,8 @@ type langfuseImpl struct {
 	restClient             *resty.Client
 	eventBuffer            *eventBuffer
 	isSendingEventsStarted atomic.Bool
+	isShutdown             atomic.Bool
+	cancelFunc             context.CancelFunc
 	endpoint               string
 	promptLabel            string
 }
@@ -42,14 +47,39 @@ func NewWithHttpClient(httpClient *http.Client, endpoint, publicKey, secretKey s
 }
 
 func (c *langfuseImpl) StartSendingEvents(ctx context.Context, period time.Duration) error {
+	if c.isShutdown.Load() {
+		return AlreadyShutdownErr
+	}
 	if c.isSendingEventsStarted.CompareAndSwap(false, true) {
+		ctx, cancel := context.WithCancel(ctx)
+		c.cancelFunc = cancel
 		go c.eventBuffer.Start(ctx, period)
 		return nil
 	} else {
 		return AlreadyStartedErr
 	}
 }
 
+func (c *langfuseImpl) Shutdown(ctx context.Context) error {
+	// Check if already shutdown
+	if c.isShutdown.CompareAndSwap(false, true) {
+		// Check if events were started
+		if !c.isSendingEventsStarted.Load() {
+			return nil
+		}
+
+		// Cancel the event buffer goroutine
+		c.cancelFunc()
+
+		// Flush any remaining events before shutdown
+		c.eventBuffer.Flush(ctx)
+
+		return nil
+	} else {
+		return AlreadyShutdownErr
+	}
+}
+
 func (c *langfuseImpl) GetSystemPromptTemplate(ctx context.Context, promptName string) (string, error) {
 	promptObject := ChatPrompt{}
 	resp, err := c.restClient.R().

langfuse_test.go

@@ -209,6 +209,92 @@ func (s *ClientTest) getClientWithMockedHttpTransport(transport httpmock.RoundTr
 	)
 }
 
+func (s *ClientTest) TestShutdownShouldReturnErrorWhenAlreadyShutdown() {
+	// Given
+	c := s.getClient()
+	err := c.StartSendingEvents(s.ctx, 1*time.Microsecond)
+	s.Require().NoError(err)
+	err = c.Shutdown(s.ctx)
+	s.Require().NoError(err)
+
+	// When
+	err = c.Shutdown(s.ctx)
+
+	// Then
+	s.Require().ErrorContains(err, "already shutdown")
+}
+
+func (s *ClientTest) TestShutdownShouldSucceedAfterStart() {
+	// Given
+	c := s.getClient()
+	err := c.StartSendingEvents(s.ctx, 1*time.Microsecond)
+	s.Require().NoError(err)
+
+	// When
+	err = c.Shutdown(s.ctx)
+
+	// Then
+	s.Require().NoError(err)
+}
+
+func (s *ClientTest) TestShutdownShouldFlushPendingEvents() {
+	// Given
+	var sentRequestBody []byte
+	wg := &sync.WaitGroup{}
+	wg.Add(1)
+	c := s.getClientWithMockedHttpTransport(func(req *http.Request) (*http.Response, error) {
+		defer wg.Done()
+		body, err := io.ReadAll(req.Body)
+		if err != nil {
+			return nil, err
+		}
+		sentRequestBody = body
+		return &http.Response{
+			StatusCode: http.StatusMultiStatus,
+			Body:       io.NopCloser(strings.NewReader("{}")),
+		}, nil
+	})
+	err := c.StartSendingEvents(s.ctx, 1*time.Hour) // Long period to prevent automatic Flush
+	s.Require().NoError(err)
+
+	// When
+	c.Trace("input", "output")
+	err = c.Shutdown(s.ctx)
+	s.Require().NoError(err)
+	wg.Wait()
+
+	// Then
+	type requestBody struct {
+		Batch []struct {
+			Body struct {
+				Input  string `json:"input"`
+				Output string `json:"output"`
+			}
+		} `json:"batch"`
+	}
+	bodyObj := requestBody{}
+	err = json.Unmarshal(sentRequestBody, &bodyObj)
+	s.Require().NoError(err)
+	s.Require().Len(bodyObj.Batch, 1)
+	s.Require().Equal("input", bodyObj.Batch[0].Body.Input)
+	s.Require().Equal("output", bodyObj.Batch[0].Body.Output)
+}
+
+func (s *ClientTest) TestStartSendingEventsShouldReturnErrorAfterShutdown() {
+	// Given
+	c := s.getClient()
+	err := c.StartSendingEvents(s.ctx, 1*time.Microsecond)
+	s.Require().NoError(err)
+	err = c.Shutdown(s.ctx)
+	s.Require().NoError(err)
+
+	// When
+	err = c.StartSendingEvents(s.ctx, 1*time.Microsecond)
+
+	// Then
+	s.Require().ErrorContains(err, "already shutdown")
+}
+
 func TestLangfuseClient(t *testing.T) {
 	suite.Run(t, new(ClientTest))
 }