From 53aa6f8a0c729eeab20e9eb821e292652de4b64d Mon Sep 17 00:00:00 2001
From: Arsen Musayelyan <moussaelianarsen@gmail.com>
Date: Fri, 22 Oct 2021 21:26:33 -0700
Subject: [PATCH] Add cancel functions to battery level and heart rate

---
 infinitime.go | 58 +++++++++++++++++++++++++++++++++------------------
 1 file changed, 38 insertions(+), 20 deletions(-)

diff --git a/infinitime.go b/infinitime.go
index 84b4ca0..798c77f 100644
--- a/infinitime.go
+++ b/infinitime.go
@@ -383,70 +383,88 @@ func (i *Device) HeartRate() (uint8, error) {
 	return uint8(heartRate[1]), nil
 }
 
-func (i *Device) WatchHeartRate() (<-chan uint8, error) {
+func (i *Device) WatchHeartRate() (<-chan uint8, func(), error) {
 	if !i.device.Properties.Connected {
-		return make(<-chan uint8), nil
+		return make(<-chan uint8), nil, nil
 	}
 	// Start notifications on heart rate characteristic
 	err := i.heartRateChar.StartNotify()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	// Watch characteristics of heart rate characteristic
 	ch, err := i.heartRateChar.WatchProperties()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	out := make(chan uint8, 2)
 	currentHeartRate, err := i.HeartRate()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	out <- currentHeartRate
+	cancel, done := cancelFunc()
 	go func() {
 		// For every event
 		for event := range ch {
-			// If value changed
-			if event.Name == "Value" {
-				// Send heart rate to channel
-				out <- uint8(event.Value.([]byte)[1])
+			select {
+			case <-done:
+				close(out)
+				close(done)
+				i.heartRateChar.StopNotify()
+				return
+			default:
+				// If value changed
+				if event.Name == "Value" {
+					// Send heart rate to channel
+					out <- uint8(event.Value.([]byte)[1])
+				}
 			}
 		}
 	}()
-	return out, nil
+	return out, cancel, nil
 }
 
-func (i *Device) WatchBatteryLevel() (<-chan uint8, error) {
+func (i *Device) WatchBatteryLevel() (<-chan uint8, func(), error) {
 	if !i.device.Properties.Connected {
-		return make(<-chan uint8), nil
+		return make(<-chan uint8), nil, nil
 	}
 	// Start notifications on heart rate characteristic
 	err := i.battLevelChar.StartNotify()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	// Watch characteristics of heart rate characteristic
 	ch, err := i.battLevelChar.WatchProperties()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	out := make(chan uint8, 2)
 	currentBattLevel, err := i.BatteryLevel()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 	out <- currentBattLevel
+	cancel, done := cancelFunc()
 	go func() {
 		// For every event
 		for event := range ch {
-			// If value changed
-			if event.Name == "Value" {
-				// Send heart rate to channel
-				out <- uint8(event.Value.([]byte)[0])
+			select {
+			case <-done:
+				close(out)
+				close(done)
+				i.battLevelChar.StopNotify()
+				return
+			default:
+				// If value changed
+				if event.Name == "Value" {
+					// Send heart rate to channel
+					out <- uint8(event.Value.([]byte)[0])
+				}
 			}
 		}
 	}()
-	return out, nil
+	return out, cancel, nil
 }
 
 func (i *Device) WatchStepCount() (<-chan uint32, func(), error) {