From b10367941aa487f6ee7cf9eaaa8179d51121a133 Mon Sep 17 00:00:00 2001
From: Joshua Rich <joshua.rich@gmail.com>
Date: Sat, 21 Dec 2024 13:48:11 +1000
Subject: [PATCH] feat(hass): :sparkles: add support to allow some requests to
 be retried

- fix default request retry logic to retry requests with 429 responses
- add ability to specify a retry flag for a request, which will retry the request on any response error using a exponential backoff mechanism
- add support for specifying retry flag for sensors and events
---
 internal/hass/api/api.go         | 20 ++++++++++++++++++++
 internal/hass/config.go          |  4 ++++
 internal/hass/event/event.go     |  9 +++++++--
 internal/hass/registration.go    |  4 ++++
 internal/hass/sensor/entities.go | 17 +++++++++++++++++
 5 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/internal/hass/api/api.go b/internal/hass/api/api.go
index d19edfcdb..3e5c407a2 100644
--- a/internal/hass/api/api.go
+++ b/internal/hass/api/api.go
@@ -24,6 +24,8 @@ const (
 var (
 	client *resty.Client
 
+	// defaultRetryFunc defines how we retry requests. By default, requests are
+	// only retried when Home Assistant responds with 429.
 	defaultRetryFunc = func(r *resty.Response, _ error) bool {
 		return r.StatusCode() == http.StatusTooManyRequests
 	}
@@ -32,11 +34,15 @@ var (
 func init() {
 	client = resty.New().
 		SetTimeout(defaultTimeout).
+		SetRetryCount(3).
+		SetRetryWaitTime(5 * time.Second).
+		SetRetryMaxWaitTime(20 * time.Second).
 		AddRetryCondition(defaultRetryFunc)
 }
 
 type Request interface {
 	RequestBody() any
+	Retry() bool
 }
 
 // Authenticated represents a request that requires passing an authentication
@@ -95,6 +101,20 @@ func Send[T any](ctx context.Context, url string, details Request) (T, error) {
 		}
 	}
 
+	if details.Retry() {
+		// If request needs to be retried, retry the request on any error.
+		logging.FromContext(ctx).Debug("Will retry requests.", slog.Any("body", details))
+		requestClient = requestClient.AddRetryCondition(
+			func(r *resty.Response, err error) bool {
+				if err != nil {
+					logging.FromContext(ctx).Debug("Retrying request.", slog.Any("body", details))
+					return true
+				}
+				return false
+			},
+		)
+	}
+
 	requestClient.SetBody(details.RequestBody())
 	responseObj, err := requestClient.Post(url)
 
diff --git a/internal/hass/config.go b/internal/hass/config.go
index 4613b59bb..268e7b963 100644
--- a/internal/hass/config.go
+++ b/internal/hass/config.go
@@ -60,3 +60,7 @@ func (c *configRequest) RequestBody() any {
 		Type: "get_config",
 	}
 }
+
+func (c *configRequest) Retry() bool {
+	return false
+}
diff --git a/internal/hass/event/event.go b/internal/hass/event/event.go
index e4e899584..09fd6fa40 100644
--- a/internal/hass/event/event.go
+++ b/internal/hass/event/event.go
@@ -15,8 +15,9 @@ const (
 )
 
 type Event struct {
-	EventData any    `json:"event_data" validate:"required"`
-	EventType string `json:"event_type" validate:"required"`
+	EventData    any    `json:"event_data" validate:"required"`
+	EventType    string `json:"event_type" validate:"required"`
+	RetryRequest bool
 }
 
 func (e *Event) Validate() error {
@@ -37,3 +38,7 @@ func (e *Event) RequestBody() any {
 		Data:        e,
 	}
 }
+
+func (e *Event) Retry() bool {
+	return e.RetryRequest
+}
diff --git a/internal/hass/registration.go b/internal/hass/registration.go
index 75cf6be27..5a6444edb 100644
--- a/internal/hass/registration.go
+++ b/internal/hass/registration.go
@@ -32,6 +32,10 @@ func (r *registrationRequest) RequestBody() any {
 	return r.Device
 }
 
+func (r *registrationRequest) Retry() bool {
+	return true
+}
+
 func newRegistrationRequest(thisDevice *device.Device, token string) *registrationRequest {
 	return &registrationRequest{
 		Device: thisDevice,
diff --git a/internal/hass/sensor/entities.go b/internal/hass/sensor/entities.go
index e07682689..4fbf8f8a5 100644
--- a/internal/hass/sensor/entities.go
+++ b/internal/hass/sensor/entities.go
@@ -25,12 +25,17 @@ type Request struct {
 	RequestType string `json:"type"`
 }
 
+type RequestMetadata struct {
+	RetryRequest bool
+}
+
 type State struct {
 	Value      any              `json:"state" validate:"required"`
 	Attributes map[string]any   `json:"attributes,omitempty" validate:"omitempty"`
 	Icon       string           `json:"icon,omitempty" validate:"omitempty,startswith=mdi:"`
 	ID         string           `json:"unique_id" validate:"required"`
 	EntityType types.SensorType `json:"type" validate:"omitempty"`
+	RequestMetadata
 }
 
 func (s *State) Validate() error {
@@ -49,6 +54,10 @@ func (s *State) RequestBody() any {
 	}
 }
 
+func (s *State) Retry() bool {
+	return s.RetryRequest
+}
+
 //nolint:wrapcheck
 func (s *State) MarshalJSON() ([]byte, error) {
 	return json.Marshal(&struct {
@@ -91,6 +100,10 @@ func (e *Entity) RequestBody() any {
 	}
 }
 
+func (e *Entity) Retry() bool {
+	return e.RetryRequest
+}
+
 //nolint:wrapcheck
 func (e *Entity) MarshalJSON() ([]byte, error) {
 	return json.Marshal(&struct {
@@ -147,3 +160,7 @@ func (l *Location) RequestBody() any {
 		Data:        l,
 	}
 }
+
+func (l *Location) Retry() bool {
+	return false
+}