pkg/server: defer ctxCleanupWG.Done()

ctxCleanupWG is being waited by the main function (or tests, see
previous commit) for the server to complete cleanup.

Commit 19b2000 introduced a return statement that did not call
.Done() on the above WaitGgroup which results in the agent hanging
if the return statement is called.

Use a defer statement for .Done() which solves the above problem and,
also, makes it harder to introduce similar bugs.

fixes: 19b2000

Signed-off-by: Kornilios Kourtis <[email protected]>

pkg/server/server.go

@@ -163,6 +163,7 @@ func (s *Server) GetEventsWG(request *tetragon.GetEventsRequest, server tetragon
 		readyWG.Done()
 	}
 	s.ctxCleanupWG.Add(1)
+	defer s.ctxCleanupWG.Done()
 	for {
 		select {
 		case event := <-l.events:
@@ -199,21 +200,18 @@ func (s *Server) GetEventsWG(request *tetragon.GetEventsRequest, server tetragon
 			} else {
 				// No need to aggregate. Directly send out the response.
 				if err = server.Send(event); err != nil {
-					s.ctxCleanupWG.Done()
 					return err
 				}
 			}
 		case <-server.Context().Done():
 			if closer != nil {
 				closer.Close()
 			}
-			s.ctxCleanupWG.Done()
 			return server.Context().Err()
 		case <-s.ctx.Done():
 			if closer != nil {
 				closer.Close()
 			}
-			s.ctxCleanupWG.Done()
 			return s.ctx.Err()
 		}
 	}