Skip to content

Commit

Permalink
Merge pull request #41278 from mindula/fix-34956
Browse files Browse the repository at this point in the history
Notify users when memory is exceeding
  • Loading branch information
KavinduZoysa authored Dec 10, 2023
2 parents af00d2c + 00bebc0 commit 4455839
Show file tree
Hide file tree
Showing 8 changed files with 226 additions and 4 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,15 @@ public interface InitializationOptions {
String KEY_ENABLE_LIGHTWEIGHT_MODE = "enableLightWeightMode";

/**
* Where the client supports inlay hints.
* Whether the client supports inlay hints.
*/
String KEY_ENABLE_INLAY_HINTS = "enableInlayHints";

/**
* Whether the client supports memory usage monitor.
*/
String KEY_ENABLE_MEMORY_USAGE_MONITOR = "enableMemoryUsageMonitor";

/**
* Return if the client support bala URI scheme.
*
Expand Down Expand Up @@ -92,4 +97,11 @@ public interface InitializationOptions {
* @return True if supported, false otherwise
*/
boolean isEnableInlayHints();

/**
* Returns if the client supports memory usage monitor.
*
* @return True if supported, false otherwise
*/
boolean isEnableMemoryUsageMonitor();
}
Original file line number Diff line number Diff line change
Expand Up @@ -40,4 +40,5 @@
requires com.google.gson;
requires io.ballerina.syntaxapicallsgen;
requires io.ballerina.central.client;
requires java.management;
}
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,10 @@ public void initialized(InitializedParams params) {
ServiceTemplateGenerator.getInstance(this.serverContext);
CentralPackageDescriptorLoader.getInstance(this.serverContext)
.loadBallerinaxPackagesFromCentral(this.serverContext);

if (lsClientCapabilities.getInitializationOptions().isEnableMemoryUsageMonitor()) {
MemoryUsageMonitor.getInstance(this.serverContext).start(client);
}
}

/**
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,11 @@ private InitializationOptions parseInitializationOptions(Map<String, Object> ini
Boolean.parseBoolean(String.valueOf(inlayHintsSupport));
initializationOptions.setEnableInlayHints(enableInlayHintsSupport);

Object memoryUsageMonitor = initOptions.get(InitializationOptions.KEY_ENABLE_MEMORY_USAGE_MONITOR);
boolean enableMemoryUsageMonitor = memoryUsageMonitor != null &&
Boolean.parseBoolean(String.valueOf(memoryUsageMonitor));
initializationOptions.setEnableMemoryUsageMonitor(enableMemoryUsageMonitor);

return initializationOptions;
}

Expand Down Expand Up @@ -214,6 +219,7 @@ public static class InitializationOptionsImpl implements InitializationOptions {
private boolean enableLSLightWeightMode = false;
private boolean supportPositionalRenamePopup = false;
private boolean enableInlayHints = false;
private boolean enableMemoryUsageMonitor = false;

@Override
public boolean isBalaSchemeSupported() {
Expand Down Expand Up @@ -267,5 +273,14 @@ public boolean isEnableInlayHints() {
public void setEnableInlayHints(boolean enableInlayHints) {
this.enableInlayHints = enableInlayHints;
}

@Override
public boolean isEnableMemoryUsageMonitor() {
return enableMemoryUsageMonitor;
}

public void setEnableMemoryUsageMonitor(boolean enableMemoryUsageMonitor) {
this.enableMemoryUsageMonitor = enableMemoryUsageMonitor;
}
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,79 @@
/*
* Copyright (c) 2023, WSO2 LLC. (http://wso2.com) All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.ballerinalang.langserver;

import org.ballerinalang.langserver.commons.LanguageServerContext;
import org.eclipse.lsp4j.MessageParams;
import org.eclipse.lsp4j.MessageType;
import org.eclipse.lsp4j.services.LanguageClient;

import java.lang.management.ManagementFactory;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryUsage;

/**
* Monitors the memory usage and notifies the user if the memory usage is high.
*
* @since 2201.9.0
*/
public class MemoryUsageMonitor {

private final MemoryMXBean memoryMXBean;
public MemoryUsageMonitor(MemoryMXBean memoryMXBean) {
this.memoryMXBean = memoryMXBean;
}

public MemoryUsageMonitor() {
this.memoryMXBean = ManagementFactory.getMemoryMXBean();
}

public static final LanguageServerContext.Key<MemoryUsageMonitor> MEMORY_USAGE_MONITOR_KEY =
new LanguageServerContext.Key<>();
public static MemoryUsageMonitor getInstance(LanguageServerContext context) {
MemoryUsageMonitor memoryUsageMonitor = context.get(MEMORY_USAGE_MONITOR_KEY);
if (memoryUsageMonitor == null) {
memoryUsageMonitor = new MemoryUsageMonitor();
context.put(MEMORY_USAGE_MONITOR_KEY, memoryUsageMonitor);
}
return memoryUsageMonitor;
}

public void start(LanguageClient client) {
Thread usageMonitor = new Thread(() -> {
while (true) {
try {
MemoryUsage heapMemoryUsage = memoryMXBean.getHeapMemoryUsage();
if (heapMemoryUsage == null) {
return;
}
long usedMemory = heapMemoryUsage.getUsed();
long maxMemory = heapMemoryUsage.getMax();

if (usedMemory >= maxMemory * 0.9) {
client.showMessage(new MessageParams(MessageType.Error,
"Memory usage is high. Some features may become unresponsive. " +
"Please reload the window or increase the memory allocated for Ballerina"));
}
Thread.sleep(60000);
} catch (InterruptedException e) {
// ignore
}
}
});
usageMonitor.setDaemon(true);
usageMonitor.start();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -832,6 +832,7 @@ public static class LanguageServerBuilder {
private OutputStream outputStream;
private InitializeParams initializeParams;
private final Map<String, Object> initOptions = new HashMap<>();
private ExtendedLanguageClient client;

public LanguageServerBuilder withLanguageServer(BallerinaLanguageServer languageServer) {
this.languageServer = languageServer;
Expand All @@ -853,6 +854,11 @@ public LanguageServerBuilder withInitOption(String key, Object value) {
return this;
}

public LanguageServerBuilder withClient(ExtendedLanguageClient client) {
this.client = client;
return this;
}

public Endpoint build() {
if (languageServer == null) {
languageServer = new BallerinaLanguageServer();
Expand All @@ -866,9 +872,12 @@ public Endpoint build() {
outputStream = OutputStream.nullOutputStream();
}

Launcher<ExtendedLanguageClient> launcher = Launcher.createLauncher(this.languageServer,
ExtendedLanguageClient.class, inputStream, outputStream);
ExtendedLanguageClient client = launcher.getRemoteProxy();
if (client == null) {
Launcher<ExtendedLanguageClient> launcher = Launcher.createLauncher(this.languageServer,
ExtendedLanguageClient.class, inputStream, outputStream);
this.client = launcher.getRemoteProxy();
}

languageServer.connect(client);

if (initializeParams == null) {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
/*
* Copyright (c) 2023, WSO2 LLC. (http://wso2.com) All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.ballerinalang.langserver.memoryusagemonitor;

import org.ballerinalang.langserver.AbstractLSTest;
import org.ballerinalang.langserver.BallerinaLanguageServer;
import org.ballerinalang.langserver.MemoryUsageMonitor;
import org.ballerinalang.langserver.commons.capability.InitializationOptions;
import org.ballerinalang.langserver.commons.client.ExtendedLanguageClient;
import org.ballerinalang.langserver.commons.workspace.WorkspaceDocumentException;
import org.ballerinalang.langserver.util.TestUtil;
import org.eclipse.lsp4j.MessageParams;
import org.eclipse.lsp4j.MessageType;
import org.eclipse.lsp4j.jsonrpc.Endpoint;
import org.mockito.Mockito;
import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

import java.io.IOException;
import java.lang.management.MemoryMXBean;
import java.lang.management.MemoryUsage;

/**
* Tests {@link MemoryUsageMonitor}.
*
* @since 2201.9.0
*/
public class MemoryUsageMonitorTest extends AbstractLSTest {
private ExtendedLanguageClient mockClient;
private MemoryMXBean mockMemoryMXBean;
private BallerinaLanguageServer languageServer;
private Endpoint serviceEndpoint;

@BeforeClass
@Override
public void init() throws Exception {
MemoryUsageMonitor memoryUsageMonitor = new MemoryUsageMonitor(createMockMemoryMXBean());
this.languageServer = new BallerinaLanguageServer();
languageServer.getServerContext().put(MemoryUsageMonitor.MEMORY_USAGE_MONITOR_KEY, memoryUsageMonitor);
mockClient = Mockito.mock(ExtendedLanguageClient.class);
}

@Test
public void test() throws WorkspaceDocumentException, IOException, InterruptedException {
TestUtil.LanguageServerBuilder builder = TestUtil.newLanguageServer()
.withLanguageServer(languageServer)
.withClient(mockClient)
.withInitOption(InitializationOptions.KEY_ENABLE_MEMORY_USAGE_MONITOR, true);
this.serviceEndpoint = builder.build();

Thread.sleep(2000);

Mockito.verify(mockClient).showMessage(new MessageParams(MessageType.Error,
"Memory usage is high. Some features may become unresponsive. " +
"Please reload the window or increase the memory allocated for Ballerina")
);
}

private MemoryMXBean createMockMemoryMXBean() {
mockMemoryMXBean = Mockito.mock(MemoryMXBean.class, Mockito.withSettings().stubOnly());
MemoryUsage mockHeapMemoryUsage = new MemoryUsage(1_000_000_000, 900_000_000, 900_000_000, 1_000_000_000);
Mockito.when(mockMemoryMXBean.getHeapMemoryUsage()).thenReturn(mockHeapMemoryUsage);
return mockMemoryMXBean;
}

@AfterClass
@Override
public void cleanMocks() {
super.cleanMocks();
if (this.mockClient != null) {
Mockito.reset(this.mockClient);
this.mockClient = null;
}
if (this.mockMemoryMXBean != null) {
Mockito.reset(this.mockMemoryMXBean);
this.mockMemoryMXBean = null;
}
}

@AfterClass
@Override
public void shutDownLanguageServer() {
TestUtil.shutdownLanguageServer(this.serviceEndpoint);
this.languageServer = null;
this.serviceEndpoint = null;
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ under the License.
<package name="org.ballerinalang.langserver.lspackageloader.*"/>
<package name="org.ballerinalang.langserver.inlayhint.*"/>
<package name="org.ballerinalang.langserver.common.utils.*"/>
<package name="org.ballerinalang.langserver.memoryusagemonitor.*"/>
</packages>
<classes>
<class name="org.ballerinalang.langserver.extensions.document.SyntaxTreeByRangeTest"/>
Expand Down

0 comments on commit 4455839

Please sign in to comment.