metron-platform/metron-integration-test/src/main/java/org/apache/metron/integration/components/FluxTopologyComponent.java - metron - Git at Google

 /**
  * Licensed to the Apache Software Foundation (ASF) under one
  * or more contributor license agreements.  See the NOTICE file
  * distributed with this work for additional information
  * regarding copyright ownership.  The ASF licenses this file
  * to you 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.apache.metron.integration.components;

 import org.apache.commons.io.FileUtils;
 import org.apache.curator.RetryPolicy;
 import org.apache.curator.framework.CuratorFramework;
 import org.apache.curator.framework.CuratorFrameworkFactory;
 import org.apache.curator.retry.ExponentialBackoffRetry;
 import org.apache.metron.integration.InMemoryComponent;
 import org.apache.metron.integration.UnableToStartException;
 import org.apache.storm.Config;
 import org.apache.storm.LocalCluster;
 import org.apache.storm.flux.FluxBuilder;
 import org.apache.storm.flux.model.ExecutionContext;
 import org.apache.storm.flux.model.TopologyDef;
 import org.apache.storm.flux.parser.FluxParser;
 import org.apache.storm.generated.KillOptions;
 import org.apache.storm.generated.StormTopology;
 import org.apache.storm.thrift.TException;
 import org.apache.storm.thrift.protocol.TProtocolException;
 import org.apache.zookeeper.data.Stat;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 import java.io.*;
 import java.lang.invoke.MethodHandles;
 import java.lang.reflect.InvocationTargetException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.util.Comparator;
 import java.util.Map;
 import java.util.Properties;
 import java.util.regex.Pattern;

 import static org.junit.jupiter.api.Assertions.assertNotNull;

 public class FluxTopologyComponent implements InMemoryComponent {

   protected static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

   LocalCluster stormCluster;
   String topologyName;
   File topologyLocation;
   File templateLocation;
   Properties topologyProperties;

   public static class Builder {

     String topologyName;
     File topologyLocation;
     File templateLocation;
     Properties topologyProperties;

     public Builder withTopologyName(String name) {
       this.topologyName = name;
       return this;
     }

     public Builder withTopologyLocation(File location) {
       this.topologyLocation = location;
       return this;
     }

     public Builder withTemplateLocation(File location) {
       this.templateLocation = location;
       return this;
     }

     public Builder withTopologyProperties(Properties properties) {
       this.topologyProperties = properties;
       this.topologyProperties.put("storm.home", "target");
       return this;
     }

     public FluxTopologyComponent build() {
       return new FluxTopologyComponent(topologyName, topologyLocation, templateLocation, topologyProperties);
     }
   }

   public FluxTopologyComponent(String topologyName, File topologyLocation, File templateLocation, Properties topologyProperties) {
     this.topologyName = topologyName;
     this.topologyLocation = topologyLocation;
     this.templateLocation = templateLocation;
     this.topologyProperties = topologyProperties;
   }

   public LocalCluster getStormCluster() {
     return stormCluster;
   }

   public String getTopologyName() {
     return topologyName;
   }

   public File getTopologyLocation() {
     return topologyLocation;
   }

   public File getTemplateLocation() {
     return templateLocation;
   }

   public Properties getTopologyProperties() {
     return topologyProperties;
   }

   public String getZookeeperConnectString() {
     return "localhost:2000";
   }

   @Override
   public void start() throws UnableToStartException {
     try {
       stormCluster = new LocalCluster();
       RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
       try(CuratorFramework client = CuratorFrameworkFactory.newClient(getZookeeperConnectString(), retryPolicy)){
         client.start();
         String root = "/storm/leader-lock";
         Stat exists = client.checkExists().forPath(root);
         if(exists == null) {
           client.create().creatingParentsIfNeeded().forPath(root);
         }
       }
       catch(Exception e) {
         LOG.error("Unable to create leaderlock", e);
       }
       finally {

       }
     } catch (Exception e) {
       throw new UnableToStartException("Unable to start flux topology: " + getTopologyLocation(), e);
     }
   }

   public static void cleanupWorkerDir() {
     if(new File("logs/workers-artifacts").exists()) {
       Path rootPath = Paths.get("logs");
       Path destPath = Paths.get("target/logs");
       try {
         Files.move(rootPath, destPath);
         Files.walk(destPath)
                 .sorted(Comparator.reverseOrder())
                 .map(Path::toFile)
                 .forEach(File::delete);
       } catch (IOException e) {
         throw new IllegalStateException(e.getMessage(), e);
       }
     }
   }

   @Override
   public void stop() {
     if (stormCluster != null) {
       try {
           try {
             // Kill the topology directly instead of sitting through the wait period
             killTopology();
             stormCluster.shutdown();
           } catch (IllegalStateException ise) {
             if (!(ise.getMessage().contains("It took over") && ise.getMessage().contains("to shut down slot"))) {
               throw ise;
             }
             else {
               LOG.error("Attempting to assassinate slots");
               assassinateSlots();
               LOG.error("Storm slots didn't shut down entirely cleanly *sigh*.  " +
                       "I gave them the old one-two-skadoo and killed the slots with prejudice.  " +
                       "If tests fail, we'll have to find a better way of killing them.", ise);
             }
         }
         catch(RuntimeException re) {
           if(re.getCause() instanceof TProtocolException) {
             //let this go, it's some intermittent weirdness.
           }
           else {
             throw re;
           }
         }
       }
       catch(Throwable t) {
         LOG.error(t.getMessage(), t);
       }
       finally {
         cleanupWorkerDir();
       }
     }
   }

   @Override
   public void reset() {
     if (stormCluster != null) {
       killTopology();
     }
   }

   protected void killTopology() {
     KillOptions ko = new KillOptions();
     ko.set_wait_secs(0);
     stormCluster.killTopologyWithOpts(topologyName, ko);
     try {
       // Actually wait for it to die.
       Thread.sleep(2000);
     } catch (InterruptedException e) {
       // Do nothing
     }
   }

   public static void assassinateSlots() {
     /*
     You might be wondering why I'm not just casting to slot here, but that's because the Slot class moved locations
     and we're supporting multiple versions of storm.
      */
     LOG.error("During slot assassination, all candidate threads: {}", Thread.getAllStackTraces().keySet());
     Thread.getAllStackTraces().keySet().stream().filter(t -> t instanceof AutoCloseable && t.getName().toLowerCase().contains("slot")).forEach(t -> {
       LOG.error("Attempting to close thread: " + t + " with state: " + t.getState());
       // With extreme prejudice.  Safety doesn't matter
       try {
         t.stop();
         LOG.error("Called thread.stop() on {}. State is: {}", t.getName(), t.getState());
       } catch(Exception e) {
         // Just swallow anything arising from the threads being killed.
       }
     });
   }

   public void submitTopology() throws NoSuchMethodException, IOException, InstantiationException, TException, IllegalAccessException, InvocationTargetException, ClassNotFoundException, NoSuchFieldException {
     startTopology(getTopologyName(), getTopologyLocation(), getTemplateLocation(), getTopologyProperties());
   }

   private void startTopology(String topologyName, File topologyLoc, File templateFile, Properties properties) throws IOException, ClassNotFoundException, NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException, TException, NoSuchFieldException{
     TopologyDef topologyDef = loadYaml(topologyName, topologyLoc, templateFile, properties);
     Config conf = FluxBuilder.buildConfig(topologyDef);
     ExecutionContext context = new ExecutionContext(topologyDef, conf);
     StormTopology topology = FluxBuilder.buildTopology(context);
     assertNotNull(topology);
     topology.validate();
     try {
       stormCluster.submitTopology(topologyName, conf, topology);
     }
     catch(Exception nne) {
       try {
         Thread.sleep(2000);
       } catch (InterruptedException e) {
       }
       stormCluster.submitTopology(topologyName, conf, topology);
     }
   }

   private static TopologyDef loadYaml(String topologyName, File yamlFile, File templateFile, Properties properties) throws IOException {
     File tmpFile = File.createTempFile(topologyName, "props");
     tmpFile.deleteOnExit();
     if (templateFile != null) {
       try (Writer propWriter = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)){
         String templateContents = FileUtils.readFileToString(templateFile);
         for(Map.Entry prop: properties.entrySet()) {
           String replacePattern = String.format("{{%s}}", prop.getKey());
           templateContents = templateContents.replaceAll(Pattern.quote(replacePattern), (String) prop.getValue());
         }
         propWriter.write(templateContents);
         propWriter.flush();
         return FluxParser.parseFile(yamlFile.getAbsolutePath(), false, true, tmpFile.getAbsolutePath(), false);
       }
     } else {
       try (Writer propWriter = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)){
         properties.store(propWriter, topologyName + " properties");
         return FluxParser.parseFile(yamlFile.getAbsolutePath(), false, true, tmpFile.getAbsolutePath(), false);
       }
     }

   }
 }
	/**
	* Licensed to the Apache Software Foundation (ASF) under one
	* or more contributor license agreements. See the NOTICE file
	* distributed with this work for additional information
	* regarding copyright ownership. The ASF licenses this file
	* to you 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.apache.metron.integration.components;

	import org.apache.commons.io.FileUtils;
	import org.apache.curator.RetryPolicy;
	import org.apache.curator.framework.CuratorFramework;
	import org.apache.curator.framework.CuratorFrameworkFactory;
	import org.apache.curator.retry.ExponentialBackoffRetry;
	import org.apache.metron.integration.InMemoryComponent;
	import org.apache.metron.integration.UnableToStartException;
	import org.apache.storm.Config;
	import org.apache.storm.LocalCluster;
	import org.apache.storm.flux.FluxBuilder;
	import org.apache.storm.flux.model.ExecutionContext;
	import org.apache.storm.flux.model.TopologyDef;
	import org.apache.storm.flux.parser.FluxParser;
	import org.apache.storm.generated.KillOptions;
	import org.apache.storm.generated.StormTopology;
	import org.apache.storm.thrift.TException;
	import org.apache.storm.thrift.protocol.TProtocolException;
	import org.apache.zookeeper.data.Stat;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	import java.io.*;
	import java.lang.invoke.MethodHandles;
	import java.lang.reflect.InvocationTargetException;
	import java.nio.charset.StandardCharsets;
	import java.nio.file.Files;
	import java.nio.file.Path;
	import java.nio.file.Paths;
	import java.util.Comparator;
	import java.util.Map;
	import java.util.Properties;
	import java.util.regex.Pattern;

	import static org.junit.jupiter.api.Assertions.assertNotNull;

	public class FluxTopologyComponent implements InMemoryComponent {

	protected static final Logger LOG = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

	LocalCluster stormCluster;
	String topologyName;
	File topologyLocation;
	File templateLocation;
	Properties topologyProperties;

	public static class Builder {

	String topologyName;
	File topologyLocation;
	File templateLocation;
	Properties topologyProperties;

	public Builder withTopologyName(String name) {
	this.topologyName = name;
	return this;
	}

	public Builder withTopologyLocation(File location) {
	this.topologyLocation = location;
	return this;
	}

	public Builder withTemplateLocation(File location) {
	this.templateLocation = location;
	return this;
	}

	public Builder withTopologyProperties(Properties properties) {
	this.topologyProperties = properties;
	this.topologyProperties.put("storm.home", "target");
	return this;
	}

	public FluxTopologyComponent build() {
	return new FluxTopologyComponent(topologyName, topologyLocation, templateLocation, topologyProperties);
	}
	}

	public FluxTopologyComponent(String topologyName, File topologyLocation, File templateLocation, Properties topologyProperties) {
	this.topologyName = topologyName;
	this.topologyLocation = topologyLocation;
	this.templateLocation = templateLocation;
	this.topologyProperties = topologyProperties;
	}

	public LocalCluster getStormCluster() {
	return stormCluster;
	}

	public String getTopologyName() {
	return topologyName;
	}

	public File getTopologyLocation() {
	return topologyLocation;
	}

	public File getTemplateLocation() {
	return templateLocation;
	}

	public Properties getTopologyProperties() {
	return topologyProperties;
	}

	public String getZookeeperConnectString() {
	return "localhost:2000";
	}

	@Override
	public void start() throws UnableToStartException {
	try {
	stormCluster = new LocalCluster();
	RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
	try(CuratorFramework client = CuratorFrameworkFactory.newClient(getZookeeperConnectString(), retryPolicy)){
	client.start();
	String root = "/storm/leader-lock";
	Stat exists = client.checkExists().forPath(root);
	if(exists == null) {
	client.create().creatingParentsIfNeeded().forPath(root);
	}
	}
	catch(Exception e) {
	LOG.error("Unable to create leaderlock", e);
	}
	finally {

	}
	} catch (Exception e) {
	throw new UnableToStartException("Unable to start flux topology: " + getTopologyLocation(), e);
	}
	}

	public static void cleanupWorkerDir() {
	if(new File("logs/workers-artifacts").exists()) {
	Path rootPath = Paths.get("logs");
	Path destPath = Paths.get("target/logs");
	try {
	Files.move(rootPath, destPath);
	Files.walk(destPath)
	.sorted(Comparator.reverseOrder())
	.map(Path::toFile)
	.forEach(File::delete);
	} catch (IOException e) {
	throw new IllegalStateException(e.getMessage(), e);
	}
	}
	}

	@Override
	public void stop() {
	if (stormCluster != null) {
	try {
	try {
	// Kill the topology directly instead of sitting through the wait period
	killTopology();
	stormCluster.shutdown();
	} catch (IllegalStateException ise) {
	if (!(ise.getMessage().contains("It took over") && ise.getMessage().contains("to shut down slot"))) {
	throw ise;
	}
	else {
	LOG.error("Attempting to assassinate slots");
	assassinateSlots();
	LOG.error("Storm slots didn't shut down entirely cleanly sigh. " +
	"I gave them the old one-two-skadoo and killed the slots with prejudice. " +
	"If tests fail, we'll have to find a better way of killing them.", ise);
	}
	}
	catch(RuntimeException re) {
	if(re.getCause() instanceof TProtocolException) {
	//let this go, it's some intermittent weirdness.
	}
	else {
	throw re;
	}
	}
	}
	catch(Throwable t) {
	LOG.error(t.getMessage(), t);
	}
	finally {
	cleanupWorkerDir();
	}
	}
	}

	@Override
	public void reset() {
	if (stormCluster != null) {
	killTopology();
	}
	}

	protected void killTopology() {
	KillOptions ko = new KillOptions();
	ko.set_wait_secs(0);
	stormCluster.killTopologyWithOpts(topologyName, ko);
	try {
	// Actually wait for it to die.
	Thread.sleep(2000);
	} catch (InterruptedException e) {
	// Do nothing
	}
	}

	public static void assassinateSlots() {
	/*
	You might be wondering why I'm not just casting to slot here, but that's because the Slot class moved locations
	and we're supporting multiple versions of storm.
	*/
	LOG.error("During slot assassination, all candidate threads: {}", Thread.getAllStackTraces().keySet());
	Thread.getAllStackTraces().keySet().stream().filter(t -> t instanceof AutoCloseable && t.getName().toLowerCase().contains("slot")).forEach(t -> {
	LOG.error("Attempting to close thread: " + t + " with state: " + t.getState());
	// With extreme prejudice. Safety doesn't matter
	try {
	t.stop();
	LOG.error("Called thread.stop() on {}. State is: {}", t.getName(), t.getState());
	} catch(Exception e) {
	// Just swallow anything arising from the threads being killed.
	}
	});
	}

	public void submitTopology() throws NoSuchMethodException, IOException, InstantiationException, TException, IllegalAccessException, InvocationTargetException, ClassNotFoundException, NoSuchFieldException {
	startTopology(getTopologyName(), getTopologyLocation(), getTemplateLocation(), getTopologyProperties());
	}

	private void startTopology(String topologyName, File topologyLoc, File templateFile, Properties properties) throws IOException, ClassNotFoundException, NoSuchMethodException, InvocationTargetException, InstantiationException, IllegalAccessException, TException, NoSuchFieldException{
	TopologyDef topologyDef = loadYaml(topologyName, topologyLoc, templateFile, properties);
	Config conf = FluxBuilder.buildConfig(topologyDef);
	ExecutionContext context = new ExecutionContext(topologyDef, conf);
	StormTopology topology = FluxBuilder.buildTopology(context);
	assertNotNull(topology);
	topology.validate();
	try {
	stormCluster.submitTopology(topologyName, conf, topology);
	}
	catch(Exception nne) {
	try {
	Thread.sleep(2000);
	} catch (InterruptedException e) {
	}
	stormCluster.submitTopology(topologyName, conf, topology);
	}
	}

	private static TopologyDef loadYaml(String topologyName, File yamlFile, File templateFile, Properties properties) throws IOException {
	File tmpFile = File.createTempFile(topologyName, "props");
	tmpFile.deleteOnExit();
	if (templateFile != null) {
	try (Writer propWriter = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)){
	String templateContents = FileUtils.readFileToString(templateFile);
	for(Map.Entry prop: properties.entrySet()) {
	String replacePattern = String.format("{{%s}}", prop.getKey());
	templateContents = templateContents.replaceAll(Pattern.quote(replacePattern), (String) prop.getValue());
	}
	propWriter.write(templateContents);
	propWriter.flush();
	return FluxParser.parseFile(yamlFile.getAbsolutePath(), false, true, tmpFile.getAbsolutePath(), false);
	}
	} else {
	try (Writer propWriter = new OutputStreamWriter(new FileOutputStream(tmpFile), StandardCharsets.UTF_8)){
	properties.store(propWriter, topologyName + " properties");
	return FluxParser.parseFile(yamlFile.getAbsolutePath(), false, true, tmpFile.getAbsolutePath(), false);
	}
	}

	}
	}